Add image format support to PdfBackend

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer
2024-10-11 16:47:15 +02:00
parent d0fccb9342
commit 6efcf0a5a5
12 changed files with 110 additions and 76 deletions

View File

@@ -7,6 +7,8 @@ from docling.backend.docling_parse_backend import (
DoclingParseDocumentBackend,
DoclingParsePageBackend,
)
from docling.datamodel.base_models import InputFormat
from docling.datamodel.document import InputDocument
@pytest.fixture
@@ -14,10 +16,21 @@ def test_doc_path():
return Path("./tests/data/2206.01062.pdf")
def _get_backend(pdf_doc):
in_doc = InputDocument(
path_or_stream=pdf_doc,
format=InputFormat.PDF,
backend=DoclingParseDocumentBackend,
)
doc_backend = in_doc._backend
return doc_backend
def test_text_cell_counts():
pdf_doc = Path("./tests/data/redp5695.pdf")
doc_backend = DoclingParseDocumentBackend(pdf_doc, "123456xyz")
doc_backend = _get_backend(pdf_doc)
for page_index in range(0, doc_backend.page_count()):
last_cell_count = None
@@ -36,7 +49,7 @@ def test_text_cell_counts():
def test_get_text_from_rect(test_doc_path):
doc_backend = DoclingParseDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
page_backend: DoclingParsePageBackend = doc_backend.load_page(0)
# Get the title text of the DocLayNet paper
@@ -49,7 +62,7 @@ def test_get_text_from_rect(test_doc_path):
def test_crop_page_image(test_doc_path):
doc_backend = DoclingParseDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
page_backend: DoclingParsePageBackend = doc_backend.load_page(0)
# Crop out "Figure 1" from the DocLayNet paper
@@ -60,5 +73,5 @@ def test_crop_page_image(test_doc_path):
def test_num_pages(test_doc_path):
doc_backend = DoclingParseDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
doc_backend.page_count() == 9

View File

@@ -7,6 +7,8 @@ from docling.backend.pypdfium2_backend import (
PyPdfiumDocumentBackend,
PyPdfiumPageBackend,
)
from docling.datamodel.base_models import InputFormat
from docling.datamodel.document import InputDocument
@pytest.fixture
@@ -14,10 +16,21 @@ def test_doc_path():
return Path("./tests/data/2206.01062.pdf")
def _get_backend(pdf_doc):
in_doc = InputDocument(
path_or_stream=pdf_doc,
format=InputFormat.PDF,
backend=PyPdfiumDocumentBackend,
)
doc_backend = in_doc._backend
return doc_backend
def test_text_cell_counts():
pdf_doc = Path("./tests/data/redp5695.pdf")
doc_backend = PyPdfiumDocumentBackend(pdf_doc, "123456xyz")
doc_backend = _get_backend(pdf_doc)
for page_index in range(0, doc_backend.page_count()):
last_cell_count = None
@@ -36,7 +49,7 @@ def test_text_cell_counts():
def test_get_text_from_rect(test_doc_path):
doc_backend = PyPdfiumDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
page_backend: PyPdfiumPageBackend = doc_backend.load_page(0)
# Get the title text of the DocLayNet paper
@@ -49,7 +62,7 @@ def test_get_text_from_rect(test_doc_path):
def test_crop_page_image(test_doc_path):
doc_backend = PyPdfiumDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
page_backend: PyPdfiumPageBackend = doc_backend.load_page(0)
# Crop out "Figure 1" from the DocLayNet paper
@@ -60,5 +73,5 @@ def test_crop_page_image(test_doc_path):
def test_num_pages(test_doc_path):
doc_backend = PyPdfiumDocumentBackend(test_doc_path, "123456xyz")
doc_backend = _get_backend(test_doc_path)
doc_backend.page_count() == 9