yujunjun/unstructured
1
0
Fork 0
mirror of https://github.com/Unstructured-IO/unstructured.git synced 2025-12-28 07:33:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feat/refactor layoutelement textregion to vectorized data structure (#3881)

Browse Source 
This PR refactors the data structure for `list[LayoutElement]` and
`list[TextRegion]` used in partition pdf/image files.

- new data structure replaces a list of objects with one object with
`numpy` array to store data
- this only affects partition internal steps and it doesn't change input
or output signature of `partition` function itself, i.e., `partition`
still returns `list[Element]`
- internally `list[LayoutElement]` -> `LayoutElements`;
`list[TextRegion]` -> `TextRegions`
- current refactor stops before clean up pdfminer elements inside
inferred layout elements -> the algorithm of clean up needs to be
refactored before the data structure refactor can move forward. So
current refactor converts the array data structure into list data
structure with `element_array.as_list()` call. This is the last step
before turning `list[LayoutElement]` into `list[Element]` as return
- a future PR will update this last step so that we build
`list[Element]` from `LayoutElements` data structure instead.

The goal of this PR is to replace the data structure as much as possible
without changing underlying logic. There are a few places where the
slicing or filtering logic was simple enough to be converted into vector
data structure operations. Those are refactored to be vector based. As a
result there is some small improvements observed in ingest test. This is
likely because the vector operations cleaned up some previous
inconsistency in data types and operations.

---------

Co-authored-by: ryannikolaidis <1208590+ryannikolaidis@users.noreply.github.com>
Co-authored-by: badGarnet <badGarnet@users.noreply.github.com>
This commit is contained in:
Yao You 2025-01-23 11:11:38 -06:00 committed by GitHub
parent 8d0b68aeae
commit 8f2a719873
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 646 additions and 428 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
CHANGELOG.md
View File

@ -1,3 +1,12 @@
## 0.16.16-dev0
### Enhancements
### Features
- **Vectorize layout (inferred, extracted, and OCR) data structure** Using `np.ndarray` to store a group of layout elements or text regions instead of using a list of objects. This improves the memory efficiency and compute speed around layout merging and deduplication.
### Fixes
## 0.16.15
### Enhancements

8
Dockerfile
View File

@ -1,7 +1,7 @@
FROM quay.io/unstructured-io/base-images:wolfi-base-latest AS base
ARG PYTHON=python3.11
ARG PIP=pip3.11
ARG PIP="${PYTHON} -m pip"
USER root
@ -19,6 +19,9 @@ RUN chown -R notebook-user:notebook-user /app && \
USER notebook-user
# append PATH before pip install to avoid warning logs; it also avoids issues with packages that needs compilation during installation
ENV PATH="${PATH}:/home/notebook-user/.local/bin"
ENV TESSDATA_PREFIX=/usr/local/share/tessdata
ENV NLTK_DATA=/home/notebook-user/nltk_data
# Install Python dependencies and download required NLTK packages
@ -28,7 +31,4 @@ RUN find requirements/ -type f -name "*.txt" -exec $PIP install --no-cache-dir -
$PYTHON -c "from unstructured.partition.model_init import initialize; initialize()" && \
$PYTHON -c "from unstructured_inference.models.tables import UnstructuredTableTransformerModel; model = UnstructuredTableTransformerModel(); model.initialize('microsoft/table-transformer-structure-recognition')"
ENV PATH="${PATH}:/home/notebook-user/.local/bin"
ENV TESSDATA_PREFIX=/usr/local/share/tessdata
CMD ["/bin/bash"]

2
Makefile
View File

@ -308,7 +308,7 @@ docker-test:
$(DOCKER_IMAGE) \
bash -c "CI=$(CI) \
UNSTRUCTURED_INCLUDE_DEBUG_METADATA=$(UNSTRUCTURED_INCLUDE_DEBUG_METADATA) \
pytest $(if $(TEST_FILE),$(TEST_FILE),test_unstructured)"
python3 -m pytest $(if $(TEST_FILE),$(TEST_FILE),test_unstructured)"
.PHONY: docker-smoke-test
docker-smoke-test:

9
test_unstructured/partition/pdf_image/test_image.py
View File

@ -79,6 +79,7 @@ class MockPageLayout(layout.PageLayout):
text="Charlie Brown and the Great Pumpkin",
),
]
self.elements_array = layout.LayoutElements.from_list(self.elements)
class MockDocumentLayout(layout.DocumentLayout):
@ -254,7 +255,10 @@ def test_partition_image_with_ocr_detects_korean():
)
assert elements[0].text == "RULES AND INSTRUCTIONS"
assert elements[3].text.replace(" ", "").startswith("안녕하세요")
# FIXME (yao): revisit this lstrip after refactoring merging logics; right now on docker and
# local testing yield different results and on docker there is a "," at the start of the Korean
# text line
assert elements[3].text.replace(" ", "").lstrip(",").startswith("안녕하세요")
def test_partition_image_with_ocr_detects_korean_from_file():
@ -267,7 +271,7 @@ def test_partition_image_with_ocr_detects_korean_from_file():
)
assert elements[0].text == "RULES AND INSTRUCTIONS"
assert elements[3].text.replace(" ", "").startswith("안녕하세요")
assert elements[3].text.replace(" ", "").lstrip(",").startswith("안녕하세요")
def test_partition_image_raises_with_bad_strategy():
@ -579,6 +583,7 @@ def inference_results():
image=mock.MagicMock(format="JPEG"),
)
page.elements = [layout.LayoutElement.from_coords(0, 0, 600, 800, text="hello")]
page.elements_array = layout.LayoutElements.from_list(page.elements)
doc = layout.DocumentLayout(pages=[page])
return doc

80
test_unstructured/partition/pdf_image/test_inference_utils.py
View File

@ -1,5 +1,5 @@
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
from unstructured.documents.elements import ElementType
from unstructured.partition.pdf_image.inference_utils import (
@ -22,16 +22,72 @@ def test_merge_text_regions(mock_embedded_text_regions):
def test_build_layout_elements_from_ocr_regions(mock_embedded_text_regions):
expected = [
LayoutElement.from_coords(
x1=437.83888888888885,
y1=317.319341111111,
x2=1256.334784222222,
y2=406.9837855555556,
text="LayoutParser: A Unified Toolkit for Deep Learning Based Document Image",
type=ElementType.UNCATEGORIZED_TEXT,
),
]
expected = LayoutElements.from_list(
[
LayoutElement.from_coords(
x1=437.83888888888885,
y1=317.319341111111,
x2=1256.334784222222,
y2=406.9837855555556,
text="LayoutParser: A Unified Toolkit for Deep Learning Based Document Image",
type=ElementType.UNCATEGORIZED_TEXT,
),
]
)
elements = build_layout_elements_from_ocr_regions(mock_embedded_text_regions)
elements = build_layout_elements_from_ocr_regions(
TextRegions.from_list(mock_embedded_text_regions)
)
assert elements == expected
def test_build_layout_elements_from_ocr_regions_with_text(mock_embedded_text_regions):
text = "LayoutParser: A Unified Toolkit for Deep Learning Based Document Image"
expected = LayoutElements.from_list(
[
LayoutElement.from_coords(
x1=437.83888888888885,
y1=317.319341111111,
x2=1256.334784222222,
y2=406.9837855555556,
text=text,
type=ElementType.UNCATEGORIZED_TEXT,
),
]
)
elements = build_layout_elements_from_ocr_regions(
TextRegions.from_list(mock_embedded_text_regions),
text,
group_by_ocr_text=True,
)
assert elements == expected
def test_build_layout_elements_from_ocr_regions_with_multi_line_text(mock_embedded_text_regions):
text = "LayoutParser: \n\nA Unified Toolkit for Deep Learning Based Document Image"
elements = build_layout_elements_from_ocr_regions(
TextRegions.from_list(mock_embedded_text_regions),
text,
group_by_ocr_text=True,
)
assert elements == LayoutElements.from_list(
[
LayoutElement.from_coords(
x1=453.00277777777774,
y1=317.319341111111,
x2=711.5338541666665,
y2=358.28571222222206,
text="LayoutParser:",
type=ElementType.UNCATEGORIZED_TEXT,
),
LayoutElement.from_coords(
x1=437.83888888888885,
y1=317.319341111111,
x2=1256.334784222222,
y2=406.9837855555556,
text="A Unified Toolkit for Deep Learning Based Document Image",
type=ElementType.UNCATEGORIZED_TEXT,
),
]
)

136
test_unstructured/partition/pdf_image/test_ocr.py
View File

@ -9,15 +9,16 @@ import unstructured_pytesseract
from bs4 import BeautifulSoup, Tag
from pdf2image.exceptions import PDFPageCountError
from PIL import Image, UnidentifiedImageError
from unstructured_inference.inference.elements import EmbeddedTextRegion, TextRegion
from unstructured_inference.inference.elements import EmbeddedTextRegion, TextRegion, TextRegions
from unstructured_inference.inference.layout import DocumentLayout
from unstructured_inference.inference.layoutelement import (
LayoutElement,
LayoutElements,
)
from unstructured.documents.elements import ElementType
from unstructured.partition.pdf_image import ocr
from unstructured.partition.pdf_image.ocr import pad_element_bboxes
from unstructured.partition.pdf_image.pdf_image_utils import pad_element_bboxes
from unstructured.partition.utils.config import env_config
from unstructured.partition.utils.constants import (
Source,
@ -90,13 +91,15 @@ def test_get_ocr_layout_from_image_tesseract(monkeypatch):
ocr_agent = OCRAgentTesseract()
ocr_layout = ocr_agent.get_layout_from_image(image)
expected_layout = [
TextRegion.from_coords(10, 5, 25, 15, "Hello", source=Source.OCR_TESSERACT),
TextRegion.from_coords(20, 15, 45, 35, "World", source=Source.OCR_TESSERACT),
TextRegion.from_coords(30, 25, 65, 55, "!", source=Source.OCR_TESSERACT),
]
expected_layout = TextRegions(
element_coords=np.array([[10.0, 5, 25, 15], [20, 15, 45, 35], [30, 25, 65, 55]]),
texts=np.array(["Hello", "World", "!"]),
sources=np.array([Source.OCR_TESSERACT] * 3),
)
assert ocr_layout == expected_layout
assert ocr_layout.texts.tolist() == expected_layout.texts.tolist()
np.testing.assert_array_equal(ocr_layout.element_coords, expected_layout.element_coords)
np.testing.assert_array_equal(ocr_layout.sources, expected_layout.sources)
def mock_ocr(*args, **kwargs):
@ -147,13 +150,15 @@ def test_get_ocr_layout_from_image_paddle(monkeypatch):
ocr_layout = OCRAgentPaddle().get_layout_from_image(image)
expected_layout = [
TextRegion.from_coords(10, 5, 25, 15, "Hello", source=Source.OCR_PADDLE),
TextRegion.from_coords(20, 15, 45, 35, "World", source=Source.OCR_PADDLE),
TextRegion.from_coords(30, 25, 65, 55, "!", source=Source.OCR_PADDLE),
]
expected_layout = TextRegions(
element_coords=np.array([[10.0, 5, 25, 15], [20, 15, 45, 35], [30, 25, 65, 55]]),
texts=np.array(["Hello", "World", "!"]),
sources=np.array([Source.OCR_PADDLE] * 3),
)
assert ocr_layout == expected_layout
assert ocr_layout.texts.tolist() == expected_layout.texts.tolist()
np.testing.assert_array_equal(ocr_layout.element_coords, expected_layout.element_coords)
np.testing.assert_array_equal(ocr_layout.sources, expected_layout.sources)
def test_get_ocr_text_from_image_tesseract(monkeypatch):
@ -254,12 +259,12 @@ def test_get_layout_from_image_google_vision(google_vision_client):
ocr_agent = google_vision_client
regions = ocr_agent.get_layout_from_image(image)
assert len(regions) == 1
assert regions[0].text == "Hello World!"
assert regions[0].source == Source.OCR_GOOGLEVISION
assert regions[0].bbox.x1 == 0
assert regions[0].bbox.y1 == 0
assert regions[0].bbox.x2 == 10
assert regions[0].bbox.y2 == 10
assert regions.texts[0] == "Hello World!"
assert all(source == Source.OCR_GOOGLEVISION for source in regions.sources)
assert regions.x1[0] == 0
assert regions.y1[0] == 0
assert regions.x2[0] == 10
assert regions.y2[0] == 10
def test_get_layout_elements_from_image_google_vision(google_vision_client):
@ -272,24 +277,28 @@ def test_get_layout_elements_from_image_google_vision(google_vision_client):
@pytest.fixture()
def mock_ocr_regions():
return [
EmbeddedTextRegion.from_coords(10, 10, 90, 90, text="0", source=None),
EmbeddedTextRegion.from_coords(200, 200, 300, 300, text="1", source=None),
EmbeddedTextRegion.from_coords(500, 320, 600, 350, text="3", source=None),
]
return TextRegions.from_list(
[
EmbeddedTextRegion.from_coords(10, 10, 90, 90, text="0", source=None),
EmbeddedTextRegion.from_coords(200, 200, 300, 300, text="1", source=None),
EmbeddedTextRegion.from_coords(500, 320, 600, 350, text="3", source=None),
]
)
@pytest.fixture()
def mock_out_layout(mock_embedded_text_regions):
return [
LayoutElement(
text=None,
source=None,
type="Text",
bbox=r.bbox,
)
for r in mock_embedded_text_regions
]
return LayoutElements.from_list(
[
LayoutElement(
text="",
source=None,
type="Text",
bbox=r.bbox,
)
for r in mock_embedded_text_regions
]
)
def test_aggregate_ocr_text_by_block():
@ -320,29 +329,31 @@ def test_zoom_image(zoom):
@pytest.fixture()
def mock_layout(mock_embedded_text_regions):
return [
LayoutElement(text=r.text, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
for r in mock_embedded_text_regions
]
return LayoutElements.from_list(
[
LayoutElement(text=r.text, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
for r in mock_embedded_text_regions
]
)
def test_supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions):
ocr_elements = [
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
for r in mock_ocr_regions
for r in mock_ocr_regions.as_list()
]
final_layout = ocr.supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions)
final_layout = ocr.supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions).as_list()
# Check if the final layout contains the original layout elements
for element in mock_layout:
for element in mock_layout.as_list():
assert element in final_layout
# Check if the final layout contains the OCR-derived elements
assert any(ocr_element in final_layout for ocr_element in ocr_elements)
# Check if the OCR-derived elements that are subregions of layout elements are removed
for element in mock_layout:
for element in mock_layout.as_list():
for ocr_element in ocr_elements:
if ocr_element.bbox.is_almost_subregion_of(
element.bbox,
@ -354,16 +365,22 @@ def test_supplement_layout_with_ocr_elements(mock_layout, mock_ocr_regions):
def test_merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions):
ocr_elements = [
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
for r in mock_ocr_regions
for r in mock_ocr_regions.as_list()
]
input_layout_elements = mock_out_layout.as_list()
final_layout = ocr.merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions)
final_layout = ocr.merge_out_layout_with_ocr_layout(
mock_out_layout,
mock_ocr_regions,
).as_list()
# Check if the out layout's text attribute is updated with aggregated OCR text
assert final_layout[0].text == mock_ocr_regions[2].text
assert final_layout[0].text == mock_ocr_regions.texts[2]
# Check if the final layout contains both original elements and OCR-derived elements
assert all(element in final_layout for element in mock_out_layout)
# The first element's text is modified by the ocr regions so it won't be the same as the input
assert all(element in final_layout for element in input_layout_elements[1:])
assert final_layout[0].bbox == input_layout_elements[0].bbox
assert any(element in final_layout for element in ocr_elements)
@ -411,11 +428,12 @@ def table_element():
@pytest.fixture()
def mock_ocr_layout():
ocr_regions = [
TextRegion.from_coords(x1=15, y1=25, x2=35, y2=45, text="Token1"),
TextRegion.from_coords(x1=40, y1=30, x2=45, y2=50, text="Token2"),
]
return ocr_regions
return TextRegions.from_list(
[
TextRegion.from_coords(x1=15, y1=25, x2=35, y2=45, text="Token1"),
TextRegion.from_coords(x1=40, y1=30, x2=45, y2=50, text="Token2"),
]
)
def test_get_table_tokens(mock_ocr_layout):
@ -462,7 +480,7 @@ def test_auto_zoom_not_exceed_tesseract_limit(monkeypatch):
image = Image.new("RGB", (1000, 1000))
ocr_agent = OCRAgentTesseract()
# tests that the code can run instead of oom and OCR results make sense
assert [region.text for region in ocr_agent.get_layout_from_image(image)] == [
assert ocr_agent.get_layout_from_image(image).texts.tolist() == [
"Hello",
"World",
"!",
@ -471,19 +489,23 @@ def test_auto_zoom_not_exceed_tesseract_limit(monkeypatch):
def test_merge_out_layout_with_cid_code(mock_out_layout, mock_ocr_regions):
# the code should ignore this invalid text and use ocr region's text
mock_out_layout[0].text = "(cid:10)(cid:5)?"
mock_out_layout.texts = mock_out_layout.texts.astype(object)
mock_out_layout.texts[0] = "(cid:10)(cid:5)?"
ocr_elements = [
LayoutElement(text=r.text, source=None, type=ElementType.UNCATEGORIZED_TEXT, bbox=r.bbox)
for r in mock_ocr_regions
for r in mock_ocr_regions.as_list()
]
input_layout_elements = mock_out_layout.as_list()
final_layout = ocr.merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions)
# TODO (yao): refactor the tests to check the array data structure directly instead of
# converting them into lists first (this includes other tests in this file)
final_layout = ocr.merge_out_layout_with_ocr_layout(mock_out_layout, mock_ocr_regions).as_list()
# Check if the out layout's text attribute is updated with aggregated OCR text
assert final_layout[0].text == mock_ocr_regions[2].text
assert final_layout[0].text == mock_ocr_regions.texts[2]
# Check if the final layout contains both original elements and OCR-derived elements
assert all(element in final_layout for element in mock_out_layout)
assert all(element in final_layout for element in input_layout_elements[1:])
assert any(element in final_layout for element in ocr_elements)

46
test_unstructured/partition/pdf_image/test_pdf.py
View File

@ -15,6 +15,7 @@ from pdf2image.exceptions import PDFPageCountError
from PIL import Image
from pytest_mock import MockFixture
from unstructured_inference.inference import layout
from unstructured_inference.inference.elements import Rectangle
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
from unstructured_inference.inference.layoutelement import LayoutElement
@ -89,22 +90,26 @@ class MockPageLayout(layout.PageLayout):
def __init__(self, number: int, image: Image):
self.number = number
self.image = image
self.image_metadata = {"width": 10, "height": 10}
self.detection_model = None
self.elements = [
layout.LayoutElement.from_coords(
type="Title",
x1=0,
y1=0,
x2=2,
y2=2,
x1=0.0,
y1=0.0,
x2=2.0,
y2=2.0,
text="Charlie Brown and the Great Pumpkin",
),
]
self.elements_array = layout.LayoutElements.from_list(self.elements)
class MockSinglePageLayout(layout.PageLayout):
def __init__(self, number: int, image: Image.Image):
self.number = number
self.image = image
self.image_metadata = {"width": 10, "height": 10}
@property
def elements(self):
@ -112,25 +117,29 @@ class MockSinglePageLayout(layout.PageLayout):
LayoutElement(
type="Headline",
text="Charlie Brown and the Great Pumpkin",
bbox=None,
bbox=Rectangle(None, None, None, None),
),
LayoutElement(
type="Subheadline",
text="The Beginning",
bbox=None,
bbox=Rectangle(None, None, None, None),
),
LayoutElement(
type="Text",
text="This time Charlie Brown had it really tricky...",
bbox=None,
bbox=Rectangle(None, None, None, None),
),
LayoutElement(
type="Title",
text="Another book title in the same page",
bbox=None,
bbox=Rectangle(None, None, None, None),
),
]
@property
def elements_array(self):
return layout.LayoutElements.from_list(self.elements)
class MockDocumentLayout(layout.DocumentLayout):
@property
@ -265,7 +274,7 @@ def test_partition_pdf_with_model_name_env_var(
with mock.patch.object(
layout,
"process_file_with_model",
mock.MagicMock(),
return_value=MockDocumentLayout(),
) as mock_process:
pdf.partition_pdf(filename=filename, strategy=PartitionStrategy.HI_RES)
assert mock_process.call_args[1]["model_name"] == "checkbox"
@ -281,7 +290,7 @@ def test_partition_pdf_with_model_name(
with mock.patch.object(
layout,
"process_file_with_model",
mock.MagicMock(),
return_value=MockDocumentLayout(),
) as mock_process:
pdf.partition_pdf(
filename=filename,
@ -293,7 +302,7 @@ def test_partition_pdf_with_model_name(
with mock.patch.object(
layout,
"process_data_with_model",
mock.MagicMock(),
return_value=MockDocumentLayout(),
) as mock_process:
with open(filename, "rb") as f:
pdf.partition_pdf(
@ -312,7 +321,7 @@ def test_partition_pdf_with_hi_res_model_name(
with mock.patch.object(
layout,
"process_file_with_model",
mock.MagicMock(),
return_value=MockDocumentLayout(),
) as mock_process:
pdf.partition_pdf(
filename=filename, strategy=PartitionStrategy.HI_RES, hi_res_model_name="checkbox"
@ -329,7 +338,7 @@ def test_partition_pdf_or_image_with_hi_res_model_name(
with mock.patch.object(
layout,
"process_file_with_model",
mock.MagicMock(),
return_value=MockDocumentLayout(),
) as mock_process:
pdf.partition_pdf_or_image(
filename=filename, strategy=PartitionStrategy.HI_RES, hi_res_model_name="checkbox"
@ -615,7 +624,9 @@ def test_partition_pdf_with_copy_protection():
def test_partition_pdf_with_dpi():
filename = example_doc_path("pdf/copy-protected.pdf")
with mock.patch.object(layout, "process_file_with_model", mock.MagicMock()) as mock_process:
with mock.patch.object(
layout, "process_file_with_model", return_value=MockDocumentLayout()
) as mock_process:
pdf.partition_pdf(filename=filename, strategy=PartitionStrategy.HI_RES, pdf_image_dpi=100)
assert mock_process.call_args[1]["pdf_image_dpi"] == 100
@ -1448,6 +1459,8 @@ def test_pdf_hi_res_max_pages_argument(filename, pdf_hi_res_max_pages, expected_
def test_document_to_element_list_omits_coord_system_when_coord_points_absent():
# TODO (yao): investigate why we need this test. The LayoutElement definition suggests bbox
# can't be None and it has to be a Rectangle object that has x1, y1, x2, y2 attributes.
layout_elem_absent_coordinates = MockSinglePageDocumentLayout()
for page in layout_elem_absent_coordinates.pages:
for el in page.elements:
@ -1463,6 +1476,7 @@ class MockImage:
format = "JPG"
@pytest.mark.skip(reason="no current layout model supports parent assignment")
def test_document_to_element_list_handles_parent():
block1 = LayoutElement.from_coords(1, 2, 3, 4, text="block 1", type="NarrativeText")
block2 = LayoutElement.from_coords(
@ -1478,7 +1492,7 @@ def test_document_to_element_list_handles_parent():
number=1,
image=MockImage(),
)
page.elements = [block1, block2]
page.elements_array = layout.LayoutElements.from_list([block1, block2])
doc = DocumentLayout.from_pages([page])
el1, el2 = pdf.document_to_element_list(doc)
assert el2.metadata.parent_id == el1.id
@ -1503,7 +1517,7 @@ def test_document_to_element_list_doesnt_sort_on_sort_method(sort_mode, call_cou
number=1,
image=MockImage(),
)
page.elements = [block1, block2]
page.elements_array = layout.LayoutElements.from_list([block1, block2])
doc = DocumentLayout.from_pages([page])
with mock.patch.object(pdf, "sort_page_elements") as mock_sort_page_elements:
pdf.document_to_element_list(doc, sortable=True, sort_mode=sort_mode)

65
test_unstructured/partition/pdf_image/test_pdfminer_processing.py
View File

@ -2,14 +2,22 @@ import numpy as np
import pytest
from PIL import Image
from unstructured_inference.constants import Source as InferenceSource
from unstructured_inference.inference.elements import EmbeddedTextRegion, Rectangle, TextRegion
from unstructured_inference.inference.elements import (
EmbeddedTextRegion,
Rectangle,
TextRegion,
TextRegions,
)
from unstructured_inference.inference.layout import DocumentLayout, LayoutElement, PageLayout
from test_unstructured.unit_utils import example_doc_path
from unstructured.partition.pdf_image.pdfminer_processing import (
_validate_bbox,
aggregate_embedded_text_by_block,
bboxes1_is_almost_subregion_of_bboxes2,
boxes_self_iou,
clean_pdfminer_inner_elements,
process_file_with_pdfminer,
remove_duplicate_elements,
)
from unstructured.partition.utils.constants import Source
@ -70,6 +78,21 @@ mix_elements_inside_table = [
]
@pytest.mark.parametrize(
("bbox", "is_valid"),
[
([0, 1, 0, 1], False),
([0, 1, 1, 2], True),
([0, 1, 1, None], False),
([0, 1, 1, np.nan], False),
([0, 1, -1, 0], False),
([0, 1, -1, 2], False),
],
)
def test_valid_bbox(bbox, is_valid):
assert _validate_bbox(bbox) is is_valid
@pytest.mark.parametrize(
("elements", "length_extra_info", "expected_document_length"),
[
@ -130,12 +153,15 @@ elements_without_duplicate_images = [
def test_aggregate_by_block():
expected = "Inside region1 Inside region2"
embedded_regions = [
TextRegion.from_coords(0, 0, 20, 20, "Inside region1"),
TextRegion.from_coords(50, 50, 150, 150, "Inside region2"),
TextRegion.from_coords(250, 250, 350, 350, "Outside region"),
]
target_region = TextRegion.from_coords(0, 0, 300, 300)
embedded_regions = TextRegions.from_list(
[
TextRegion.from_coords(0, 0, 20, 20, "Inside region1"),
TextRegion.from_coords(20, 20, 80, 80, None),
TextRegion.from_coords(50, 50, 150, 150, "Inside region2"),
TextRegion.from_coords(250, 250, 350, 350, "Outside region"),
]
)
target_region = TextRegions.from_list([TextRegion.from_coords(0, 0, 300, 300)])
text = aggregate_embedded_text_by_block(target_region, embedded_regions)
assert text == expected
@ -195,19 +221,24 @@ def test_boxes_self_iou(coords, threshold, expected):
def test_remove_duplicate_elements():
sample_elements = [
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 1"),
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 2"),
EmbeddedTextRegion(bbox=Rectangle(20, 20, 30, 30), text="Text 3"),
]
sample_elements = TextRegions.from_list(
[
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 1"),
EmbeddedTextRegion(bbox=Rectangle(0, 0, 10, 10), text="Text 2"),
EmbeddedTextRegion(bbox=Rectangle(20, 20, 30, 30), text="Text 3"),
]
)
result = remove_duplicate_elements(sample_elements)
# Check that duplicates were removed and only 2 unique elements remain
assert len(result) == 2
assert result[0].text == "Text 2"
assert result[1].text == "Text 3"
assert result.texts.tolist() == ["Text 2", "Text 3"]
assert result.element_coords.tolist() == [[0, 0, 10, 10], [20, 20, 30, 30]]
# Ensure the duplicate was removed by checking that result contains no redundant bboxes
assert result[0].bbox == Rectangle(0, 0, 10, 10)
assert result[1].bbox == Rectangle(20, 20, 30, 30)
def test_process_file_with_pdfminer():
layout, links = process_file_with_pdfminer(example_doc_path("pdf/layout-parser-paper-fast.pdf"))
assert len(layout)
assert "LayoutParser: A Unified Toolkit for Deep\n" in layout[0].texts
assert links[0][0]["url"] == "https://layout-parser.github.io"

30
test_unstructured/partition/utils/test_sorting.py
View File

@ -1,4 +1,6 @@
import numpy as np
import pytest
from unstructured_inference.inference.elements import TextRegions
from unstructured.documents.coordinates import PixelSpace
from unstructured.documents.elements import CoordinatesMetadata, Element, Text
@ -8,6 +10,7 @@ from unstructured.partition.utils.sorting import (
coordinates_to_bbox,
shrink_bbox,
sort_page_elements,
sort_text_regions,
)
@ -109,6 +112,33 @@ def test_sort_basic_pos_coordinates():
assert sorted_elem_text == "7 8 9"
def test_sort_text_regions():
unsorted = TextRegions(
element_coords=np.array(
[[1, 2, 2, 2], [1, 1, 2, 2], [3, 1, 4, 4]],
),
texts=np.array(["1", "2", "3"]),
sources=np.array(["foo"] * 3),
)
assert sort_text_regions(unsorted, sort_mode=SORT_MODE_BASIC).texts.tolist() == ["2", "3", "1"]
@pytest.mark.parametrize(
"coords",
[
[[1, 2, 2, 2], [1, 1, 2, 2], [3, -1, 4, 4]],
[[1, 2, 2, 2], [1, 1, 2, 2], [3, None, 4, 4]],
],
)
def test_sort_text_regions_with_invalid_coords_using_xy_cut_does_no_ops(coords):
unsorted = TextRegions(
element_coords=np.array(coords).astype(float),
texts=np.array(["1", "2", "3"]),
sources=np.array(["foo"] * 3),
)
assert sort_text_regions(unsorted).texts.tolist() == ["1", "2", "3"]
def test_coordinates_to_bbox():
coordinates_data = MockCoordinatesMetadata([(10, 20), (10, 200), (100, 200), (100, 20)])
expected_result = (10, 20, 100, 200)

46
test_unstructured_ingest/expected-structured-output/local-single-file-with-pdf-infer-table-structure/layout-parser-paper-with-table.jpg.json
View File

@ -66,32 +66,10 @@
}
}
},
{
"type": "UncategorizedText",
"element_id": "e5314387378c7a98911d71c145c45327",
"text": "2",
"metadata": {
"filetype": "image/jpeg",
"languages": [
"eng"
],
"page_number": 1,
"data_source": {
"record_locator": {
"path": "/home/runner/work/unstructured/unstructured/test_unstructured_ingest/example-docs/layout-parser-paper-with-table.jpg"
},
"permissions_data": [
{
"mode": 33188
}
]
}
}
},
{
"type": "FigureCaption",
"element_id": "e262996994d01c45f0d6ef28cb8afa93",
"text": "For each dataset, we train several models of different sizes for different needs (the trade-off between accuracy vs. computational cost). For \u201cbase model\u201d and \u201clarge model\u201d, we refer to using the ResNet 50 or ResNet 101 backbones [13], respectively. One can train models of different architectures, like Faster R-CNN [28] (P) and Mask R-CNN [12] (M). For example, an F in the Large Model column indicates it has m Faster R-CNN model trained using the ResNet 101 backbone. The platform is maintained and a number of additions will be made to the model zoo in coming months.",
"element_id": "a0c3c6b7e1e8c95016b989ef43c5ea2e",
"text": "2 For each dataset, we train several models of different sizes for different needs (the trade-off between accuracy vs. computational cost). For \u201cbase model\u201d and \u201clarge model\u201d, we refer to using the ResNet 50 or ResNet 101 backbones [13], respectively. One can train models of different architectures, like Faster R-CNN [28] (P) and Mask R-CNN [12] (M). For example, an F in the Large Model column indicates it has m Faster R-CNN model trained using the ResNet 101 backbone. The platform is maintained and a number of additions will be made to the model zoo in coming months.",
"metadata": {
"filetype": "image/jpeg",
"languages": [
@ -112,7 +90,7 @@
},
{
"type": "NarrativeText",
"element_id": "2298258fe84201e839939d70c168141b",
"element_id": "b68ca269882f83b03827b5edf0fec979",
"text": "layout data structures, which are optimized for efficiency and versatility. 3) When necessary, users can employ existing or customized OCR models via the unified API provided in the OCR module. 4) LayoutParser comes with a set of utility functions for the visualization and stomge of the layout data. 5) LayoutParser is also highly customizable, via its integration with functions for layout data annotation and model training. We now provide detailed descriptions for each component.",
"metadata": {
"filetype": "image/jpeg",
@ -134,7 +112,7 @@
},
{
"type": "Title",
"element_id": "24d2473c4975fedd3f5cfd3026249837",
"element_id": "a98721b4c18e53da7ee4e38512d91480",
"text": "3.1 Layout Detection Models",
"metadata": {
"filetype": "image/jpeg",
@ -156,7 +134,7 @@
},
{
"type": "NarrativeText",
"element_id": "008c0a590378dccd98ae7a5c49905eda",
"element_id": "84bf4abf7f899f83b876d112cbe176f4",
"text": "In LayoutParser, a layout model takes a document image as an input and generates a list of rectangular boxes for the target content regions. Different from traditional methods, it relies on deep convolutional neural networks rather than manually curated rules to identify content regions. It is formulated as an object detection problem and state-of-the-art models like Faster R-CNN [28] and Mask R-CNN [12] are used. This yields prediction results of high accuracy and makes it possible to build a concise, generalized interface for layout detection. LayoutParser, built upon Detectron2 [35], provides a minimal API that can perform layout detection with only four lines of code in Python:",
"metadata": {
"filetype": "image/jpeg",
@ -178,7 +156,7 @@
},
{
"type": "ListItem",
"element_id": "b98aac79b1c1af144f6ed563e6510fd4",
"element_id": "04d62ad595016d7b490dff67a00b9f35",
"text": "import layoutparser as lp",
"metadata": {
"filetype": "image/jpeg",
@ -200,7 +178,7 @@
},
{
"type": "Title",
"element_id": "44691a14713d40ea25a0401490ed7b5e",
"element_id": "9d40bf1b2e2af1692f5689a1c44ab2ae",
"text": "wwe",
"metadata": {
"filetype": "image/jpeg",
@ -222,7 +200,7 @@
},
{
"type": "ListItem",
"element_id": "e14922762abe8a044371efcab13bdcc9",
"element_id": "cafbdebf75706654ed769cd9785e8697",
"text": "image = cv2.imread(\"image_file\") # load images",
"metadata": {
"filetype": "image/jpeg",
@ -244,7 +222,7 @@
},
{
"type": "ListItem",
"element_id": "986e6a00c43302413ca0ad4badd5bca8",
"element_id": "e8455ed7a816cc15906468871b66a90a",
"text": "model = lp. Detectron2LayoutModel (",
"metadata": {
"filetype": "image/jpeg",
@ -266,7 +244,7 @@
},
{
"type": "ListItem",
"element_id": "d50233678a0d15373eb47ab537d3c11e",
"element_id": "44fd87fd2c9870a523e3b8cc3483da53",
"text": "ea \"lp: //PubLayNet/faster_rcnn_R_50_FPN_3x/config\")",
"metadata": {
"filetype": "image/jpeg",
@ -288,7 +266,7 @@
},
{
"type": "ListItem",
"element_id": "11dccdd53ee27c94e976b875d2d6e40d",
"element_id": "f4db9091ab6b62feee72d2bde0ff9e87",
"text": "layout = model.detect (image)",
"metadata": {
"filetype": "image/jpeg",
@ -310,7 +288,7 @@
},
{
"type": "NarrativeText",
"element_id": "bb86a9374cb6126db4088d1092557d09",
"element_id": "e277edc46744590708425e453eea87c1",
"text": "LayoutParser provides a wealth of pre-trained model weights using various datasets covering different languages, time periods, and document types. Due to domain shift [7], the prediction performance can notably drop when models are ap- plied to target samples that are significantly different from the training dataset. As document structures and layouts vary greatly in different domains, it is important to select models trained on a dataset similar to the test samples. A semantic syntax is used for initializing the model weights in Layout Parser, using both the dataset name and model name 1p://<dataset-name>/<model-architecture-name>.",
"metadata": {
"filetype": "image/jpeg",

2
unstructured/__version__.py
View File

@ -1 +1 @@
__version__ = "0.16.15" # pragma: no cover
__version__ = "0.16.16-dev0" # pragma: no cover

2
unstructured/partition/common/common.py
View File

@ -53,7 +53,7 @@ def normalize_layout_element(
text = layout_dict.get("text", "")
# Both `coordinates` and `coordinate_system` must be present
# in order to add coordinates metadata to the element.
coordinates = layout_dict.get("coordinates")
coordinates = layout_dict.get("coordinates") if coordinate_system else None
element_type = layout_dict.get("type")
prob = layout_dict.get("prob")
aux_origin = layout_dict.get("source", None)

15
unstructured/partition/pdf.py
View File

@ -613,7 +613,7 @@ def _partition_pdf_or_image_local(
model_name=hi_res_model_name,
)
extracted_layout_dumper = ExtractedLayoutDumper(
layout=extracted_layout,
layout=[layout.as_list() for layout in extracted_layout],
)
ocr_layout_dumper = OCRLayoutDumper()
# NOTE(christine): merged_document_layout = extracted_layout + inferred_layout
@ -665,7 +665,7 @@ def _partition_pdf_or_image_local(
model_name=hi_res_model_name,
)
extracted_layout_dumper = ExtractedLayoutDumper(
layout=extracted_layout,
layout=[layout.as_list() for layout in extracted_layout],
)
ocr_layout_dumper = OCRLayoutDumper()
@ -690,6 +690,7 @@ def _partition_pdf_or_image_local(
ocr_layout_dumper=ocr_layout_dumper,
)
# vectorization of the data structure ends here
final_document_layout = clean_pdfminer_inner_elements(final_document_layout)
for page in final_document_layout.pages:
@ -903,8 +904,10 @@ def _partition_pdf_or_image_with_ocr_from_image(
languages=languages,
)
# NOTE (yao): elements for a document is still stored as a list therefore at this step we have
# to convert the vector data structured ocr_data into a list
page_elements = ocr_data_to_elements(
ocr_data,
ocr_data.as_list(),
image_size=image.size,
common_metadata=metadata,
)
@ -1123,7 +1126,11 @@ def document_to_element_list(
)
for layout_element in page.elements:
if image_width and image_height and hasattr(layout_element.bbox, "coordinates"):
if (
image_width
and image_height
and getattr(layout_element.bbox, "x1") not in (None, np.nan)
):
coordinate_system = PixelSpace(width=image_width, height=image_height)
else:
coordinate_system = None

44
unstructured/partition/pdf_image/inference_utils.py
View File

@ -2,10 +2,12 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import numpy as np
from unstructured_inference.constants import Source
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layoutelement import (
LayoutElement,
LayoutElements,
partition_groups_from_regions,
)
@ -39,44 +41,45 @@ def build_layout_element(
def build_layout_elements_from_ocr_regions(
ocr_regions: list[TextRegion],
ocr_regions: TextRegions,
ocr_text: Optional[str] = None,
group_by_ocr_text: bool = False,
) -> list[LayoutElement]:
) -> LayoutElements:
"""
Get layout elements from OCR regions
"""
grouped_regions = []
if group_by_ocr_text:
text_sections = ocr_text.split("\n\n")
grouped_regions = []
mask = np.ones(ocr_regions.texts.shape).astype(bool)
indices = np.arange(len(mask))
for text_section in text_sections:
regions = []
words = text_section.replace("\n", " ").split()
for ocr_region in ocr_regions:
for i, text in enumerate(ocr_regions.texts[mask]):
if not words:
break
if ocr_region.text in words:
regions.append(ocr_region)
words.remove(ocr_region.text)
if text in words:
regions.append(indices[mask][i])
mask[mask][i] = False
words.remove(text)
if not regions:
continue
for r in regions:
ocr_regions.remove(r)
grouped_regions.append(TextRegions.from_list(regions))
grouped_regions.append(ocr_regions.slice(regions))
else:
grouped_regions = partition_groups_from_regions(TextRegions.from_list(ocr_regions))
grouped_regions = partition_groups_from_regions(ocr_regions)
merged_regions = [merge_text_regions(group) for group in grouped_regions]
return [
build_layout_element(
bbox=r.bbox, text=r.text, source=r.source, element_type=ElementType.UNCATEGORIZED_TEXT
)
for r in merged_regions
]
merged_regions = TextRegions.from_list([merge_text_regions(group) for group in grouped_regions])
return LayoutElements(
element_coords=merged_regions.element_coords,
texts=merged_regions.texts,
sources=merged_regions.sources,
element_class_ids=np.zeros(merged_regions.texts.shape),
element_class_id_map={0: ElementType.UNCATEGORIZED_TEXT},
)
def merge_text_regions(regions: TextRegions) -> TextRegion:
@ -99,6 +102,7 @@ def merge_text_regions(regions: TextRegions) -> TextRegion:
max_y2 = regions.y2.max().astype(float)
merged_text = " ".join([text for text in regions.texts if text])
source = regions.source
# assumption is the regions has the same source
source = regions.sources[0]
return TextRegion.from_coords(min_x1, min_y1, max_x2, max_y2, merged_text, source)

181
unstructured/partition/pdf_image/ocr.py
View File

@ -4,6 +4,7 @@ import os
import tempfile
from typing import IO, TYPE_CHECKING, Any, List, Optional, cast
import numpy as np
import pdf2image
# NOTE(yuming): Rename PIL.Image to avoid conflict with
@ -14,16 +15,20 @@ from PIL import ImageSequence
from unstructured.documents.elements import ElementType
from unstructured.metrics.table.table_formats import SimpleTableCell
from unstructured.partition.pdf_image.analysis.layout_dump import OCRLayoutDumper
from unstructured.partition.pdf_image.pdf_image_utils import pad_element_bboxes, valid_text
from unstructured.partition.pdf_image.pdf_image_utils import valid_text
from unstructured.partition.pdf_image.pdfminer_processing import (
aggregate_embedded_text_by_block,
bboxes1_is_almost_subregion_of_bboxes2,
)
from unstructured.partition.utils.config import env_config
from unstructured.partition.utils.constants import OCRMode
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
from unstructured.utils import requires_dependencies
if TYPE_CHECKING:
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layout import DocumentLayout, PageLayout
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
from unstructured_inference.models.tables import UnstructuredTableTransformerModel
@ -93,7 +98,7 @@ def process_data_with_ocr(
def process_file_with_ocr(
filename: str,
out_layout: "DocumentLayout",
extracted_layout: List[List["TextRegion"]],
extracted_layout: List[TextRegions],
is_image: bool = False,
infer_table_structure: bool = False,
ocr_languages: str = "eng",
@ -110,6 +115,9 @@ def process_file_with_ocr(
- out_layout (DocumentLayout): The output layout from unstructured-inference.
- extracted_layout (List[TextRegions]): a list of text regions extracted by pdfminer, one for
each page
- is_image (bool, optional): Indicates if the input data is an image (True) or not (False).
Defaults to False.
@ -187,7 +195,7 @@ def supplement_page_layout_with_ocr(
infer_table_structure: bool = False,
ocr_languages: str = "eng",
ocr_mode: str = OCRMode.FULL_PAGE.value,
extracted_regions: Optional[List["TextRegion"]] = None,
extracted_regions: Optional[TextRegions] = None,
ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
) -> "PageLayout":
"""
@ -202,28 +210,30 @@ def supplement_page_layout_with_ocr(
if ocr_mode == OCRMode.FULL_PAGE.value:
ocr_layout = ocr_agent.get_layout_from_image(image)
if ocr_layout_dumper:
ocr_layout_dumper.add_ocred_page(ocr_layout)
page_layout.elements[:] = merge_out_layout_with_ocr_layout(
out_layout=cast(List["LayoutElement"], page_layout.elements),
ocr_layout_dumper.add_ocred_page(ocr_layout.as_list())
page_layout.elements_array = merge_out_layout_with_ocr_layout(
out_layout=page_layout.elements_array,
ocr_layout=ocr_layout,
)
elif ocr_mode == OCRMode.INDIVIDUAL_BLOCKS.value:
for element in page_layout.elements:
if not element.text:
padding = env_config.IMAGE_CROP_PAD
padded_element = pad_element_bboxes(element, padding=padding)
cropped_image = image.crop(
(
padded_element.bbox.x1,
padded_element.bbox.y1,
padded_element.bbox.x2,
padded_element.bbox.y2,
),
)
# Note(yuming): instead of getting OCR layout, we just need
# the text extraced from OCR for individual elements
text_from_ocr = ocr_agent.get_text_from_image(cropped_image)
element.text = text_from_ocr
# individual block mode still keeps using the list data structure for elements instead of
# the vectorized page_layout.elements_array data structure
for i, text in enumerate(page_layout.elements_array.texts):
if text:
continue
padding = env_config.IMAGE_CROP_PAD
cropped_image = image.crop(
(
page_layout.elements_array.x1[i] - padding,
page_layout.elements_array.y1[i] - padding,
page_layout.elements_array.x2[i] + padding,
page_layout.elements_array.y2[i] + padding,
),
)
# Note(yuming): instead of getting OCR layout, we just need
# the text extraced from OCR for individual elements
text_from_ocr = ocr_agent.get_text_from_image(cropped_image)
page_layout.elements_array.texts[i] = text_from_ocr
else:
raise ValueError(
"Invalid OCR mode. Parameter `ocr_mode` "
@ -238,24 +248,25 @@ def supplement_page_layout_with_ocr(
if tables.tables_agent is None:
raise RuntimeError("Unable to load table extraction agent.")
page_layout.elements[:] = supplement_element_with_table_extraction(
elements=cast(List["LayoutElement"], page_layout.elements),
page_layout.elements_array = supplement_element_with_table_extraction(
elements=page_layout.elements_array,
image=image,
tables_agent=tables.tables_agent,
ocr_agent=ocr_agent,
extracted_regions=extracted_regions,
)
page_layout.elements = page_layout.elements_array.as_list()
return page_layout
@requires_dependencies("unstructured_inference")
def supplement_element_with_table_extraction(
elements: List["LayoutElement"],
elements: LayoutElements,
image: PILImage.Image,
tables_agent: "UnstructuredTableTransformerModel",
ocr_agent,
extracted_regions: Optional[List["TextRegion"]] = None,
extracted_regions: Optional[TextRegions] = None,
) -> List["LayoutElement"]:
"""Supplement the existing layout with table extraction. Any Table elements
that are extracted will have a metadata fields "text_as_html" where
@ -264,23 +275,26 @@ def supplement_element_with_table_extraction(
"""
from unstructured_inference.models.tables import cells_to_html
table_elements = [el for el in elements if el.type == ElementType.TABLE]
for element in table_elements:
padding = env_config.TABLE_IMAGE_CROP_PAD
padded_element = pad_element_bboxes(element, padding=padding)
table_id = {v: k for k, v in elements.element_class_id_map.items()}.get(ElementType.TABLE)
if not table_id:
# no table found in this page
return elements
table_ele_indices = np.where(elements.element_class_ids == table_id)[0]
table_elements = elements.slice(table_ele_indices)
padding = env_config.TABLE_IMAGE_CROP_PAD
for i, element_coords in enumerate(table_elements.element_coords):
cropped_image = image.crop(
(
padded_element.bbox.x1,
padded_element.bbox.y1,
padded_element.bbox.x2,
padded_element.bbox.y2,
element_coords[0] - padding,
element_coords[1] - padding,
element_coords[2] + padding,
element_coords[3] + padding,
),
)
table_tokens = get_table_tokens(
table_element_image=cropped_image,
ocr_agent=ocr_agent,
extracted_regions=extracted_regions,
table_element=padded_element,
)
tatr_cells = tables_agent.predict(
cropped_image, ocr_tokens=table_tokens, result_format="cells"
@ -288,13 +302,13 @@ def supplement_element_with_table_extraction(
# NOTE(christine): `tatr_cells == ""` means that the table was not recognized
text_as_html = "" if tatr_cells == "" else cells_to_html(tatr_cells)
element.text_as_html = text_as_html
elements.text_as_html[table_ele_indices[i]] = text_as_html
if env_config.EXTRACT_TABLE_AS_CELLS:
simple_table_cells = [
SimpleTableCell.from_table_transformer_cell(cell).to_dict() for cell in tatr_cells
]
element.table_as_cells = simple_table_cells
elements.table_as_cells[table_ele_indices[i]] = simple_table_cells
return elements
@ -302,44 +316,38 @@ def supplement_element_with_table_extraction(
def get_table_tokens(
table_element_image: PILImage.Image,
ocr_agent: OCRAgent,
extracted_regions: Optional[List["TextRegion"]] = None,
table_element: Optional["LayoutElement"] = None,
) -> List[dict[str, Any]]:
"""Get OCR tokens from either paddleocr or tesseract"""
ocr_layout = ocr_agent.get_layout_from_image(image=table_element_image)
table_tokens = []
for ocr_region in ocr_layout:
for i, text in enumerate(ocr_layout.texts):
table_tokens.append(
{
"bbox": [
ocr_region.bbox.x1,
ocr_region.bbox.y1,
ocr_region.bbox.x2,
ocr_region.bbox.y2,
ocr_layout.x1[i],
ocr_layout.y1[i],
ocr_layout.x2[i],
ocr_layout.y2[i],
],
"text": ocr_region.text,
"text": text,
# 'table_tokens' is a list of tokens
# Need to be in a relative reading order
"span_num": i,
"line_num": 0,
"block_num": 0,
}
)
# 'table_tokens' is a list of tokens
# Need to be in a relative reading order
# If no order is provided, use current order
for idx, token in enumerate(table_tokens):
if "span_num" not in token:
token["span_num"] = idx
if "line_num" not in token:
token["line_num"] = 0
if "block_num" not in token:
token["block_num"] = 0
return table_tokens
def merge_out_layout_with_ocr_layout(
out_layout: List["LayoutElement"],
ocr_layout: List["TextRegion"],
out_layout: LayoutElements,
ocr_layout: TextRegions,
supplement_with_ocr_elements: bool = True,
) -> List["LayoutElement"]:
subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
) -> LayoutElements:
"""
Merge the out layout with the OCR-detected text regions on page level.
@ -349,12 +357,14 @@ def merge_out_layout_with_ocr_layout(
supplemented with the OCR layout.
"""
out_regions_without_text = [region for region in out_layout if not valid_text(region.text)]
invalid_text_indices = [i for i, text in enumerate(out_layout.texts) if not valid_text(text)]
out_layout.texts = out_layout.texts.astype(object)
for out_region in out_regions_without_text:
out_region.text = aggregate_ocr_text_by_block(
ocr_layout,
out_region,
for idx in invalid_text_indices:
out_layout.texts[idx] = aggregate_embedded_text_by_block(
target_region=out_layout.slice([idx]),
source_regions=ocr_layout,
threshold=subregion_threshold,
)
final_layout = (
@ -389,10 +399,10 @@ def aggregate_ocr_text_by_block(
@requires_dependencies("unstructured_inference")
def supplement_layout_with_ocr_elements(
layout: List["LayoutElement"],
ocr_layout: List["TextRegion"],
layout: LayoutElements,
ocr_layout: TextRegions,
subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
) -> List["LayoutElement"]:
) -> LayoutElements:
"""
Supplement the existing layout with additional OCR-derived elements.
@ -402,10 +412,8 @@ def supplement_layout_with_ocr_elements(
OCR-derived list. Then, it appends the remaining OCR-derived regions to the existing layout.
Parameters:
- layout (List[LayoutElement]): A list of existing layout elements, each of which is
an instance of `LayoutElement`.
- ocr_layout (List[TextRegion]): A list of OCR-derived text regions, each of which is
an instance of `TextRegion`.
- layout (LayoutElements): A collection of existing layout elements in array structures
- ocr_layout (TextRegions): A collection of OCR-derived text regions in array structures
Returns:
- List[LayoutElement]: The final combined layout consisting of both the original layout
@ -420,25 +428,26 @@ def supplement_layout_with_ocr_elements(
threshold.
"""
from unstructured_inference.inference.layoutelement import LayoutElements
from unstructured.partition.pdf_image.inference_utils import (
build_layout_elements_from_ocr_regions,
)
ocr_regions_to_remove: list[TextRegion] = []
for ocr_region in ocr_layout:
for el in layout:
ocr_region_is_subregion_of_out_el = ocr_region.bbox.is_almost_subregion_of(
el.bbox,
subregion_threshold,
)
if ocr_region_is_subregion_of_out_el:
ocr_regions_to_remove.append(ocr_region)
break
mask = (
~bboxes1_is_almost_subregion_of_bboxes2(
ocr_layout.element_coords, layout.element_coords, subregion_threshold
)
.sum(axis=1)
.astype(bool)
)
ocr_regions_to_add = [region for region in ocr_layout if region not in ocr_regions_to_remove]
if ocr_regions_to_add:
# add ocr regions that are not covered by layout
ocr_regions_to_add = ocr_layout.slice(mask)
if sum(mask):
ocr_elements_to_add = build_layout_elements_from_ocr_regions(ocr_regions_to_add)
final_layout = layout + ocr_elements_to_add
final_layout = LayoutElements.concatenate([layout, ocr_elements_to_add])
else:
final_layout = layout

244
unstructured/partition/pdf_image/pdfminer_processing.py
View File

@ -23,8 +23,9 @@ from unstructured.partition.utils.sorting import sort_text_regions
from unstructured.utils import requires_dependencies
if TYPE_CHECKING:
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layout import DocumentLayout
from unstructured_inference.inference.layoutelement import LayoutElements
EPSILON_AREA = 0.01
@ -45,18 +46,79 @@ def process_file_with_pdfminer(
return extracted_layout, layouts_links
def _validate_bbox(bbox: list[int | float]) -> bool:
return all(x is not None for x in bbox) and (bbox[2] - bbox[0] > 0) and (bbox[3] - bbox[1] > 0)
@requires_dependencies("unstructured_inference")
def process_page_layout_from_pdfminer(
annotation_list: list,
page_layout,
page_height: int | float,
page_number: int,
coord_coef: float,
) -> tuple[LayoutElements, list]:
from unstructured_inference.inference.layoutelement import LayoutElements
urls_metadata: list[dict[str, Any]] = []
element_coords, texts, element_class = [], [], []
annotation_threshold = env_config.PDF_ANNOTATION_THRESHOLD
for obj in page_layout:
x1, y1, x2, y2 = rect_to_bbox(obj.bbox, page_height)
bbox = (x1, y1, x2, y2)
if len(annotation_list) > 0 and isinstance(obj, LTTextBox):
annotations_within_element = check_annotations_within_element(
annotation_list,
bbox,
page_number,
annotation_threshold,
)
_, words = get_words_from_obj(obj, page_height)
for annot in annotations_within_element:
urls_metadata.append(map_bbox_and_index(words, annot))
if hasattr(obj, "get_text"):
inner_text_objects = extract_text_objects(obj)
for inner_obj in inner_text_objects:
inner_bbox = rect_to_bbox(inner_obj.bbox, page_height)
if not _validate_bbox(inner_bbox):
continue
texts.append(inner_obj.get_text())
element_coords.append(inner_bbox)
element_class.append(0)
else:
inner_image_objects = extract_image_objects(obj)
for img_obj in inner_image_objects:
inner_bbox = rect_to_bbox(img_obj.bbox, page_height)
if not _validate_bbox(inner_bbox):
continue
texts.append(None)
element_coords.append(inner_bbox)
element_class.append(1)
return (
LayoutElements(
element_coords=coord_coef * np.array(element_coords),
texts=np.array(texts).astype(object),
element_class_ids=np.array(element_class),
element_class_id_map={0: "Text", 1: "Image"},
sources=np.array([Source.PDFMINER] * len(element_class)),
),
urls_metadata,
)
@requires_dependencies("unstructured_inference")
def process_data_with_pdfminer(
file: Optional[Union[bytes, BinaryIO]] = None,
dpi: int = 200,
) -> tuple[List[List["TextRegion"]], List[List]]:
) -> tuple[List[LayoutElements], List[List]]:
"""Loads the image and word objects from a pdf using pdfplumber and the image renderings of the
pdf pages using pdf2image"""
from unstructured_inference.inference.elements import (
EmbeddedTextRegion,
ImageTextRegion,
)
from unstructured_inference.inference.layoutelement import LayoutElements
layouts = []
layouts_links = []
@ -65,8 +127,6 @@ def process_data_with_pdfminer(
for page_number, (page, page_layout) in enumerate(open_pdfminer_pages_generator(file)):
width, height = page_layout.width, page_layout.height
text_layout = []
image_layout = []
annotation_list = []
coordinate_system = PixelSpace(
width=width,
@ -75,49 +135,10 @@ def process_data_with_pdfminer(
if page.annots:
annotation_list = get_uris(page.annots, height, coordinate_system, page_number)
annotation_threshold = env_config.PDF_ANNOTATION_THRESHOLD
urls_metadata: list[dict[str, Any]] = []
layout, urls_metadata = process_page_layout_from_pdfminer(
annotation_list, page_layout, height, page_number, coef
)
for obj in page_layout:
x1, y1, x2, y2 = rect_to_bbox(obj.bbox, height)
bbox = (x1, y1, x2, y2)
if len(annotation_list) > 0 and isinstance(obj, LTTextBox):
annotations_within_element = check_annotations_within_element(
annotation_list,
bbox,
page_number,
annotation_threshold,
)
_, words = get_words_from_obj(obj, height)
for annot in annotations_within_element:
urls_metadata.append(map_bbox_and_index(words, annot))
if hasattr(obj, "get_text"):
inner_text_objects = extract_text_objects(obj)
for inner_obj in inner_text_objects:
_text = inner_obj.get_text()
text_region = _create_text_region(
*rect_to_bbox(inner_obj.bbox, height),
coef,
_text,
Source.PDFMINER,
EmbeddedTextRegion,
)
if text_region.bbox is not None and text_region.bbox.area > 0:
text_layout.append(text_region)
else:
inner_image_objects = extract_image_objects(obj)
for img_obj in inner_image_objects:
text_region = _create_text_region(
*rect_to_bbox(img_obj.bbox, height),
coef,
None,
Source.PDFMINER,
ImageTextRegion,
)
if text_region.bbox is not None and text_region.bbox.area > 0:
image_layout.append(text_region)
links = [
{
"bbox": [x * coef for x in metadata["bbox"]],
@ -128,13 +149,22 @@ def process_data_with_pdfminer(
for metadata in urls_metadata
]
clean_text_layout = remove_duplicate_elements(
text_layout, env_config.EMBEDDED_TEXT_SAME_REGION_THRESHOLD
)
clean_image_layout = remove_duplicate_elements(
image_layout, env_config.EMBEDDED_IMAGE_SAME_REGION_THRESHOLD
)
layout = [*clean_text_layout, *clean_image_layout]
clean_layouts = []
for threshold, element_class in zip(
(
env_config.EMBEDDED_TEXT_SAME_REGION_THRESHOLD,
env_config.EMBEDDED_IMAGE_SAME_REGION_THRESHOLD,
),
(0, 1),
):
elements_to_sort = layout.slice(layout.element_class_ids == element_class)
clean_layouts.append(
remove_duplicate_elements(elements_to_sort, threshold)
if len(elements_to_sort)
else elements_to_sort
)
layout = LayoutElements.concatenate(clean_layouts)
# NOTE(christine): always do the basic sort first for deterministic order across
# python versions.
layout = sort_text_regions(layout, SORT_MODE_BASIC)
@ -161,6 +191,9 @@ def _create_text_region(x1, y1, x2, y2, coef, text, source, region_class):
def get_coords_from_bboxes(bboxes, round_to: int = DEFAULT_ROUND) -> np.ndarray:
"""convert a list of boxes's coords into np array"""
if isinstance(bboxes, np.ndarray):
return bboxes.round(round_to)
# preallocate memory
coords = np.zeros((len(bboxes), 4), dtype=np.float32)
@ -214,14 +247,38 @@ def boxes_self_iou(bboxes, threshold: float = 0.5, round_to: int = DEFAULT_ROUND
return (inter_area / np.maximum(EPSILON_AREA, boxa_area + boxb_area.T - inter_area)) > threshold
@requires_dependencies("unstructured_inference")
def pdfminer_elements_to_text_regions(layout_elements: LayoutElements) -> list[TextRegions]:
"""a temporary solution to convert layout elements to a list of either EmbeddedTextRegion or
ImageTextRegion; this should be made obsolete after we refactor the merging logic in inference
library"""
from unstructured_inference.inference.elements import (
EmbeddedTextRegion,
ImageTextRegion,
)
regions = []
for i, element_class in enumerate(layout_elements.element_class_ids):
region_class = EmbeddedTextRegion if element_class == 0 else ImageTextRegion
regions.append(
region_class.from_coords(
*layout_elements.element_coords[i],
text=layout_elements.texts[i],
source=Source.PDFMINER,
)
)
return regions
@requires_dependencies("unstructured_inference")
def merge_inferred_with_extracted_layout(
inferred_document_layout: "DocumentLayout",
extracted_layout: List[List["TextRegion"]],
extracted_layout: List[TextRegions],
hi_res_model_name: str,
) -> "DocumentLayout":
"""Merge an inferred layout with an extracted layout"""
from unstructured_inference.inference.layoutelement import LayoutElements
from unstructured_inference.inference.layoutelement import (
merge_inferred_layout_with_extracted_layout as merge_inferred_with_extracted_page,
)
@ -246,28 +303,30 @@ def merge_inferred_with_extracted_layout(
):
threshold_kwargs = {"same_region_threshold": 0.5, "subregion_threshold": 0.5}
# NOTE (yao): after refactoring the algorithm to be vectorized we can then pass in the
# vectorized data structure into the merge function
merged_layout = merge_inferred_with_extracted_page(
inferred_layout=inferred_layout,
extracted_layout=extracted_page_layout,
extracted_layout=pdfminer_elements_to_text_regions(extracted_page_layout),
page_image_size=image_size,
**threshold_kwargs,
)
merged_layout = sort_text_regions(cast(List["TextRegion"], merged_layout), SORT_MODE_BASIC)
merged_layout = sort_text_regions(LayoutElements.from_list(merged_layout), SORT_MODE_BASIC)
# so that we can modify the text without worrying about hitting length limit
merged_layout.texts = merged_layout.texts.astype(object)
elements = []
for layout_el in merged_layout:
if layout_el.text is None:
for i, text in enumerate(merged_layout.texts):
if text is None:
text = aggregate_embedded_text_by_block(
text_region=cast("TextRegion", layout_el),
pdf_objects=extracted_page_layout,
target_region=merged_layout.slice([i]),
source_regions=extracted_page_layout,
)
else:
text = layout_el.text
layout_el.text = remove_control_characters(text)
elements.append(layout_el)
merged_layout.texts[i] = remove_control_characters(text)
inferred_page.elements[:] = elements
inferred_page.elements_array = merged_layout
# NOTE: once we drop reference to elements we can remove this step below
inferred_page.elements[:] = merged_layout.as_list()
return inferred_document_layout
@ -313,40 +372,39 @@ def clean_pdfminer_inner_elements(document: "DocumentLayout") -> "DocumentLayout
@requires_dependencies("unstructured_inference")
def remove_duplicate_elements(
elements: list["TextRegion"],
elements: TextRegions,
threshold: float = 0.5,
) -> list["TextRegion"]:
) -> TextRegions:
"""Removes duplicate text elements extracted by PDFMiner from a document layout."""
bboxes = []
for i, element in enumerate(elements):
bboxes.append(element.bbox)
iou = boxes_self_iou(bboxes, threshold)
filtered_elements = []
for i, element in enumerate(elements):
if iou[i, i + 1 :].any():
continue
filtered_elements.append(element)
return filtered_elements
iou = boxes_self_iou(elements.element_coords, threshold)
# this is equivalent of finding those rows where `not iou[i, i + 1 :].any()`, i.e., any element
# that has no overlap above the threshold with any other elements
return elements.slice(~np.triu(iou, k=1).any(axis=1))
def aggregate_embedded_text_by_block(
text_region: "TextRegion",
pdf_objects: list["TextRegion"],
target_region: TextRegions,
source_regions: TextRegions,
threshold: float = env_config.EMBEDDED_TEXT_AGGREGATION_SUBREGION_THRESHOLD,
) -> str:
"""Extracts the text aggregated from the elements of the given layout that lie within the given
block."""
mask = bboxes1_is_almost_subregion_of_bboxes2(
[obj.bbox for obj in pdf_objects],
[text_region.bbox],
env_config.EMBEDDED_TEXT_AGGREGATION_SUBREGION_THRESHOLD,
).sum(axis=1)
if len(source_regions) == 0 or len(target_region) == 0:
return ""
text = " ".join([obj.text for i, obj in enumerate(pdf_objects) if (mask[i] and obj.text)])
mask = (
bboxes1_is_almost_subregion_of_bboxes2(
source_regions.element_coords,
target_region.element_coords,
threshold,
)
.sum(axis=1)
.astype(bool)
)
text = " ".join([text for text in source_regions.slice(mask).texts if text])
return text

18
unstructured/partition/utils/ocr_models/google_vision_ocr.py
View File

@ -12,8 +12,8 @@ from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
if TYPE_CHECKING:
from PIL import Image as PILImage
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layoutelement import LayoutElements
class OCRAgentGoogleVision(OCRAgent):
@ -44,7 +44,7 @@ class OCRAgentGoogleVision(OCRAgent):
assert isinstance(document, TextAnnotation)
return document.text
def get_layout_from_image(self, image: PILImage.Image) -> list[TextRegion]:
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
trace_logger.detail("Processing entire page OCR with Google Vision API...")
image_context = ImageContext(language_hints=[self.language]) if self.language else None
with BytesIO() as buffer:
@ -57,7 +57,8 @@ class OCRAgentGoogleVision(OCRAgent):
regions = self._parse_regions(document)
return regions
def get_layout_elements_from_image(self, image: PILImage.Image) -> list[LayoutElement]:
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
from unstructured.partition.pdf_image.inference_utils import (
build_layout_elements_from_ocr_regions,
)
@ -68,14 +69,15 @@ class OCRAgentGoogleVision(OCRAgent):
ocr_text = self.get_text_from_image(
image,
)
layout_elements = build_layout_elements_from_ocr_regions(
return build_layout_elements_from_ocr_regions(
ocr_regions=ocr_regions,
ocr_text=ocr_text,
group_by_ocr_text=False,
)
return layout_elements
def _parse_regions(self, ocr_data: TextAnnotation) -> list[TextRegion]:
def _parse_regions(self, ocr_data: TextAnnotation) -> TextRegions:
from unstructured_inference.inference.elements import TextRegions
from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
text_regions: list[TextRegion] = []
@ -94,7 +96,7 @@ class OCRAgentGoogleVision(OCRAgent):
source=Source.OCR_GOOGLEVISION,
)
text_regions.append(text_region)
return text_regions
return TextRegions.from_list(text_regions)
def _get_text_from_paragraph(self, paragraph: Paragraph) -> str:
breaks = TextAnnotation.DetectedBreak.BreakType

8
unstructured/partition/utils/ocr_models/ocr_interface.py
View File

@ -17,8 +17,8 @@ from unstructured.partition.utils.constants import (
if TYPE_CHECKING:
from PIL import Image as PILImage
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.elements import TextRegions
from unstructured_inference.inference.layoutelement import LayoutElements
class OCRAgent(ABC):
@ -55,11 +55,11 @@ class OCRAgent(ABC):
)
@abstractmethod
def get_layout_elements_from_image(self, image: PILImage.Image) -> list[LayoutElement]:
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
pass
@abstractmethod
def get_layout_from_image(self, image: PILImage.Image) -> list[TextRegion]:
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
pass
@abstractmethod

39
unstructured/partition/utils/ocr_models/paddle_ocr.py
View File

@ -12,8 +12,8 @@ from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
from unstructured.utils import requires_dependencies
if TYPE_CHECKING:
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.elements import TextRegion, TextRegions
from unstructured_inference.inference.layoutelement import LayoutElements
class OCRAgentPaddle(OCRAgent):
@ -61,12 +61,12 @@ class OCRAgentPaddle(OCRAgent):
def get_text_from_image(self, image: PILImage.Image) -> str:
ocr_regions = self.get_layout_from_image(image)
return "\n\n".join([r.text for r in ocr_regions])
return "\n\n".join(ocr_regions.texts)
def is_text_sorted(self):
return False
def get_layout_from_image(self, image: PILImage.Image) -> list[TextRegion]:
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
"""Get the OCR regions from image as a list of text regions with paddle."""
trace_logger.detail("Processing entire page OCR with paddle...")
@ -80,26 +80,22 @@ class OCRAgentPaddle(OCRAgent):
return ocr_regions
@requires_dependencies("unstructured_inference")
def get_layout_elements_from_image(self, image: PILImage.Image) -> list[LayoutElement]:
from unstructured.partition.pdf_image.inference_utils import build_layout_element
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
ocr_regions = self.get_layout_from_image(image)
# NOTE(christine): For paddle, there is no difference in `ocr_layout` and `ocr_text` in
# terms of grouping because we get ocr_text from `ocr_layout, so the first two grouping
# and merging steps are not necessary.
return [
build_layout_element(
bbox=r.bbox,
text=r.text,
source=r.source,
element_type=ElementType.UNCATEGORIZED_TEXT,
)
for r in ocr_regions
]
return LayoutElements(
element_coords=ocr_regions.element_coords,
texts=ocr_regions.texts,
element_class_ids=np.zeros(ocr_regions.texts.shape),
element_class_id_map={0: ElementType.UNCATEGORIZED_TEXT},
)
@requires_dependencies("unstructured_inference")
def parse_data(self, ocr_data: list[Any]) -> list[TextRegion]:
def parse_data(self, ocr_data: list[Any]) -> TextRegions:
"""Parse the OCR result data to extract a list of TextRegion objects from paddle.
The function processes the OCR result dictionary, looking for bounding
@ -110,14 +106,17 @@ class OCRAgentPaddle(OCRAgent):
- ocr_data (list): A list containing the OCR result data
Returns:
- list[TextRegion]: A list of TextRegion objects, each representing a
detected text region within the OCR-ed image.
- TextRegions:
TextRegions object, containing data from all text regions in numpy arrays; each row
represents a detected text region within the OCR-ed image.
Note:
- An empty string or a None value for the 'text' key in the input
dictionary will result in its associated bounding box being ignored.
"""
from unstructured_inference.inference.elements import TextRegions
from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
text_regions: list[TextRegion] = []
@ -141,4 +140,6 @@ class OCRAgentPaddle(OCRAgent):
)
text_regions.append(text_region)
return text_regions
# FIXME (yao): find out if paddle supports a vectorized output format so we can skip the
# step of parsing a list
return TextRegions.from_list(text_regions)

52
unstructured/partition/utils/ocr_models/tesseract_ocr.py
View File

@ -2,7 +2,7 @@ from __future__ import annotations
import os
import re
from typing import TYPE_CHECKING, List
from typing import TYPE_CHECKING
import cv2
import numpy as np
@ -23,8 +23,8 @@ from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
from unstructured.utils import requires_dependencies
if TYPE_CHECKING:
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.layoutelement import LayoutElement
from unstructured_inference.inference.elements import TextRegions
from unstructured_inference.inference.layoutelement import LayoutElements
# -- force tesseract to be single threaded, otherwise we see major performance problems --
if "OMP_THREAD_LIMIT" not in os.environ:
@ -43,7 +43,7 @@ class OCRAgentTesseract(OCRAgent):
def get_text_from_image(self, image: PILImage.Image) -> str:
return unstructured_pytesseract.image_to_string(np.array(image), lang=self.language)
def get_layout_from_image(self, image: PILImage.Image) -> List[TextRegion]:
def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
"""Get the OCR regions from image as a list of text regions with tesseract."""
trace_logger.detail("Processing entire page OCR with tesseract...")
@ -166,7 +166,7 @@ class OCRAgentTesseract(OCRAgent):
return word_text
@requires_dependencies("unstructured_inference")
def get_layout_elements_from_image(self, image: PILImage.Image) -> List["LayoutElement"]:
def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
from unstructured.partition.pdf_image.inference_utils import (
build_layout_elements_from_ocr_regions,
)
@ -189,7 +189,7 @@ class OCRAgentTesseract(OCRAgent):
)
@requires_dependencies("unstructured_inference")
def parse_data(self, ocr_data: pd.DataFrame, zoom: float = 1) -> List["TextRegion"]:
def parse_data(self, ocr_data: pd.DataFrame, zoom: float = 1) -> TextRegions:
"""Parse the OCR result data to extract a list of TextRegion objects from tesseract.
The function processes the OCR result data frame, looking for bounding
@ -206,39 +206,33 @@ class OCRAgentTesseract(OCRAgent):
Default is 1.
Returns:
- List[TextRegion]:
A list of TextRegion objects, each representing a detected text region
within the OCR-ed image.
- TextRegions:
TextRegions object, containing data from all text regions in numpy arrays; each row
represents a detected text region within the OCR-ed image.
Note:
- An empty string or a None value for the 'text' key in the input
data frame will result in its associated bounding box being ignored.
"""
from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
from unstructured_inference.inference.elements import TextRegions
if zoom <= 0:
zoom = 1
text_regions: list[TextRegion] = []
for idtx in ocr_data.itertuples():
text = idtx.text
if not text:
continue
cleaned_text = str(text) if not isinstance(text, str) else text.strip()
if cleaned_text:
x1 = idtx.left / zoom
y1 = idtx.top / zoom
x2 = (idtx.left + idtx.width) / zoom
y2 = (idtx.top + idtx.height) / zoom
text_region = build_text_region_from_coords(
x1, y1, x2, y2, text=cleaned_text, source=Source.OCR_TESSERACT
)
text_regions.append(text_region)
return text_regions
texts = ocr_data.text.apply(
lambda text: str(text) if not isinstance(text, str) else text.strip()
).values
mask = texts != ""
element_coords = ocr_data[["left", "top", "width", "height"]].values
element_coords[:, 2] += element_coords[:, 0]
element_coords[:, 3] += element_coords[:, 1]
element_coords = element_coords.astype(float) / zoom
return TextRegions(
element_coords=element_coords[mask],
texts=texts[mask],
sources=np.array([Source.OCR_TESSERACT] * mask.sum()),
)
def zoom_image(image: PILImage.Image, zoom: float = 1) -> PILImage.Image:

38
unstructured/partition/utils/sorting.py
View File

@ -11,7 +11,7 @@ from unstructured.partition.utils.constants import SORT_MODE_BASIC, SORT_MODE_XY
from unstructured.partition.utils.xycut import recursive_xy_cut, recursive_xy_cut_swapped
if TYPE_CHECKING:
from unstructured_inference.inference.elements import TextRegion
from unstructured_inference.inference.elements import TextRegions
def coordinates_to_bbox(coordinates: CoordinatesMetadata) -> tuple[int, int, int, int]:
@ -213,33 +213,30 @@ def sort_bboxes_by_xy_cut(
def sort_text_regions(
elements: list["TextRegion"],
elements: TextRegions,
sort_mode: str = SORT_MODE_XY_CUT,
shrink_factor: float = 0.9,
xy_cut_primary_direction: str = "x",
) -> list["TextRegion"]:
) -> TextRegions:
"""Sort a list of TextRegion elements based on the specified sorting mode."""
if not elements:
return elements
bboxes = [(el.bbox.x1, el.bbox.y1, el.bbox.x2, el.bbox.y2) for el in elements]
bboxes = elements.element_coords
def _bboxes_ok(strict_points: bool):
warned = False
for bbox in bboxes:
if bbox is None:
trace_logger.detail( # type: ignore
"some or all elements are missing bboxes, skipping sort",
)
if np.isnan(bboxes).any():
trace_logger.detail( # type: ignore
"some or all elements are missing bboxes, skipping sort",
)
return False
if bboxes.shape[1] != 4 or np.where(bboxes < 0)[0].size:
trace_logger.detail("at least one bbox contains invalid values") # type: ignore
if strict_points:
return False
elif not bbox_is_valid(bbox):
if not warned:
trace_logger.detail(f"bbox {bbox} does not have valid values") # type: ignore
warned = True
if strict_points:
return False
return True
if sort_mode == SORT_MODE_XY_CUT:
@ -260,11 +257,12 @@ def sort_text_regions(
shrink_factor=shrink_factor,
xy_cut_primary_direction=xy_cut_primary_direction,
)
sorted_elements = [elements[i] for i in res]
sorted_elements = elements.slice(res)
elif sort_mode == SORT_MODE_BASIC:
sorted_elements = sorted(
elements,
key=lambda el: (el.bbox.y1, el.bbox.x1, el.bbox.y2, el.bbox.x2),
# NOTE (yao): lexsort order is revese from the input sequence; so below is first sort by y1,
# then x1, then y2, lastly x2
sorted_elements = elements.slice(
np.lexsort((elements.x2, elements.y2, elements.x1, elements.y1))
)
else:
sorted_elements = elements