docling 1.19.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

docling/models/table_structure_model.py

@@ -3,29 +3,25 @@ from pathlib import Path
 from typing import Iterable, List
 
 import numpy
+from docling_core.types.doc import BoundingBox, DocItemLabel, TableCell
 from docling_ibm_models.tableformer.data_management.tf_predictor import TFPredictor
 from PIL import ImageDraw
 
-from docling.datamodel.base_models import (
-    BoundingBox,
-    Page,
-    TableCell,
-    TableElement,
-    TableStructurePrediction,
-)
-from docling.datamodel.pipeline_options import TableFormerMode
+from docling.datamodel.base_models import Page, Table, TableStructurePrediction
+from docling.datamodel.pipeline_options import TableFormerMode, TableStructureOptions
+from docling.models.base_model import BasePageModel
 
 
-class TableStructureModel:
-    def __init__(self, config):
-        self.config = config
-        self.do_cell_matching = config["do_cell_matching"]
-        self.mode = config["mode"]
+class TableStructureModel(BasePageModel):
+    def __init__(
+        self, enabled: bool, artifacts_path: Path, options: TableStructureOptions
+    ):
+        self.options = options
+        self.do_cell_matching = self.options.do_cell_matching
+        self.mode = self.options.mode
 
-        self.enabled = config["enabled"]
+        self.enabled = enabled
         if self.enabled:
-            artifacts_path: Path = config["artifacts_path"]
-
             if self.mode == TableFormerMode.ACCURATE:
                 artifacts_path = artifacts_path / "fat"
 
@@ -39,7 +35,9 @@ class TableStructureModel:
             self.tf_predictor = TFPredictor(self.tm_config)
             self.scale = 2.0  # Scale up table input images to 144 dpi
 
-    def draw_table_and_cells(self, page: Page, tbl_list: List[TableElement]):
+    def draw_table_and_cells(self, page: Page, tbl_list: List[Table]):
+        assert page._backend is not None
+
         image = (
             page._backend.get_page_image()
         )  # make new image to avoid drawing on the saved ones
@@ -50,17 +48,18 @@ class TableStructureModel:
             draw.rectangle([(x0, y0), (x1, y1)], outline="red")
 
             for tc in table_element.table_cells:
-                x0, y0, x1, y1 = tc.bbox.as_tuple()
-                if tc.column_header:
-                    width = 3
-                else:
-                    width = 1
-                draw.rectangle([(x0, y0), (x1, y1)], outline="blue", width=width)
-                draw.text(
-                    (x0 + 3, y0 + 3),
-                    text=f"{tc.start_row_offset_idx}, {tc.start_col_offset_idx}",
-                    fill="black",
-                )
+                if tc.bbox is not None:
+                    x0, y0, x1, y1 = tc.bbox.as_tuple()
+                    if tc.column_header:
+                        width = 3
+                    else:
+                        width = 1
+                    draw.rectangle([(x0, y0), (x1, y1)], outline="blue", width=width)
+                    draw.text(
+                        (x0 + 3, y0 + 3),
+                        text=f"{tc.start_row_offset_idx}, {tc.start_col_offset_idx}",
+                        fill="black",
+                    )
 
         image.show()
 
@@ -71,6 +70,9 @@ class TableStructureModel:
             return
 
         for page in page_batch:
+            assert page._backend is not None
+            assert page.predictions.layout is not None
+            assert page.size is not None
 
             page.predictions.tablestructure = TableStructurePrediction()  # dummy
 
@@ -85,7 +87,7 @@ class TableStructureModel:
                     ],
                 )
                 for cluster in page.predictions.layout.clusters
-                if cluster.label == "Table"
+                if cluster.label == DocItemLabel.TABLE
             ]
             if not len(in_tables):
                 yield page
@@ -132,7 +134,7 @@ class TableStructureModel:
                             element["bbox"]["token"] = text_piece
 
                         tc = TableCell.model_validate(element)
-                        if self.do_cell_matching:
+                        if self.do_cell_matching and tc.bbox is not None:
                             tc.bbox = tc.bbox.scaled(1 / self.scale)
                         table_cells.append(tc)
 
@@ -141,7 +143,7 @@ class TableStructureModel:
                     num_cols = table_out["predict_details"]["num_cols"]
                     otsl_seq = table_out["predict_details"]["prediction"]["rs_seq"]
 
-                    tbl = TableElement(
+                    tbl = Table(
                         otsl_seq=otsl_seq,
                         table_cells=table_cells,
                         num_rows=num_rows,
@@ -149,7 +151,7 @@ class TableStructureModel:
                         id=table_cluster.id,
                         page_no=page.page_no,
                         cluster=table_cluster,
-                        label="Table",
+                        label=DocItemLabel.TABLE,
                     )
 
                     page.predictions.tablestructure.table_map[table_cluster.id] = tbl

docling/models/tesseract_ocr_cli_model.py

@@ -2,11 +2,12 @@ import io
 import logging
 import tempfile
 from subprocess import DEVNULL, PIPE, Popen
-from typing import Iterable, Tuple
+from typing import Iterable, Optional, Tuple
 
 import pandas as pd
+from docling_core.types.doc import BoundingBox, CoordOrigin
 
-from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell, Page
+from docling.datamodel.base_models import OcrCell, Page
 from docling.datamodel.pipeline_options import TesseractCliOcrOptions
 from docling.models.base_ocr_model import BaseOcrModel
 
@@ -21,8 +22,8 @@ class TesseractOcrCliModel(BaseOcrModel):
 
         self.scale = 3  # multiplier for 72 dpi == 216 dpi.
 
-        self._name = None
-        self._version = None
+        self._name: Optional[str] = None
+        self._version: Optional[str] = None
 
         if self.enabled:
             try:
@@ -39,7 +40,7 @@ class TesseractOcrCliModel(BaseOcrModel):
     def _get_name_and_version(self) -> Tuple[str, str]:
 
         if self._name != None and self._version != None:
-            return self._name, self._version
+            return self._name, self._version  # type: ignore
 
         cmd = [self.options.tesseract_cmd, "--version"]
 
@@ -108,6 +109,8 @@ class TesseractOcrCliModel(BaseOcrModel):
             return
 
         for page in page_batch:
+            assert page._backend is not None
+
             ocr_rects = self.get_ocr_rects(page)
 
             all_ocr_cells = []

docling/models/tesseract_ocr_model.py

@@ -1,19 +1,19 @@
 import logging
 from typing import Iterable
 
-import numpy
+from docling_core.types.doc import BoundingBox, CoordOrigin
 
-from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell, Page
-from docling.datamodel.pipeline_options import TesseractCliOcrOptions
+from docling.datamodel.base_models import OcrCell, Page
+from docling.datamodel.pipeline_options import TesseractOcrOptions
 from docling.models.base_ocr_model import BaseOcrModel
 
 _log = logging.getLogger(__name__)
 
 
 class TesseractOcrModel(BaseOcrModel):
-    def __init__(self, enabled: bool, options: TesseractCliOcrOptions):
+    def __init__(self, enabled: bool, options: TesseractOcrOptions):
         super().__init__(enabled=enabled, options=options)
-        self.options: TesseractCliOcrOptions
+        self.options: TesseractOcrOptions
 
         self.scale = 3  # multiplier for 72 dpi == 216 dpi.
         self.reader = None
@@ -68,6 +68,9 @@ class TesseractOcrModel(BaseOcrModel):
             return
 
         for page in page_batch:
+            assert page._backend is not None
+            assert self.reader is not None
+
             ocr_rects = self.get_ocr_rects(page)
 
             all_ocr_cells = []

docling/pipeline/base_pipeline.py

@@ -0,0 +1,190 @@
+import functools
+import logging
+import time
+import traceback
+from abc import ABC, abstractmethod
+from typing import Callable, Iterable, List
+
+from docling_core.types.doc import DoclingDocument, NodeItem
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.backend.pdf_backend import PdfDocumentBackend
+from docling.datamodel.base_models import (
+    ConversionStatus,
+    DoclingComponentType,
+    ErrorItem,
+    Page,
+)
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.datamodel.settings import settings
+from docling.models.base_model import BaseEnrichmentModel
+from docling.utils.utils import chunkify
+
+_log = logging.getLogger(__name__)
+
+
+class BasePipeline(ABC):
+    def __init__(self, pipeline_options: PipelineOptions):
+        self.pipeline_options = pipeline_options
+        self.build_pipe: List[Callable] = []
+        self.enrichment_pipe: List[BaseEnrichmentModel] = []
+
+    def execute(self, in_doc: InputDocument, raises_on_error: bool) -> ConversionResult:
+        conv_res = ConversionResult(input=in_doc)
+
+        _log.info(f"Processing document {in_doc.file.name}")
+        try:
+            # These steps are building and assembling the structure of the
+            # output DoclingDocument
+            conv_res = self._build_document(in_doc, conv_res)
+            conv_res = self._assemble_document(in_doc, conv_res)
+            # From this stage, all operations should rely only on conv_res.output
+            conv_res = self._enrich_document(in_doc, conv_res)
+            conv_res.status = self._determine_status(in_doc, conv_res)
+        except Exception as e:
+            conv_res.status = ConversionStatus.FAILURE
+            if raises_on_error:
+                raise e
+
+        return conv_res
+
+    @abstractmethod
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        pass
+
+    def _assemble_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        return conv_res
+
+    def _enrich_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        def _filter_elements(
+            doc: DoclingDocument, model: BaseEnrichmentModel
+        ) -> Iterable[NodeItem]:
+            for element, _level in doc.iterate_items():
+                if model.is_processable(doc=doc, element=element):
+                    yield element
+
+        for model in self.enrichment_pipe:
+            for element_batch in chunkify(
+                _filter_elements(conv_res.document, model),
+                settings.perf.elements_batch_size,
+            ):
+                # TODO: currently we assume the element itself is modified, because
+                # we don't have an interface to save the element back to the document
+                for element in model(
+                    doc=conv_res.document, element_batch=element_batch
+                ):  # Must exhaust!
+                    pass
+
+        return conv_res
+
+    @abstractmethod
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_default_options(cls) -> PipelineOptions:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        pass
+
+    # def _apply_on_elements(self, element_batch: Iterable[NodeItem]) -> Iterable[Any]:
+    #    for model in self.build_pipe:
+    #        element_batch = model(element_batch)
+    #
+    #    yield from element_batch
+
+
+class PaginatedPipeline(BasePipeline):  # TODO this is a bad name.
+
+    def _apply_on_pages(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+        for model in self.build_pipe:
+            page_batch = model(page_batch)
+
+        yield from page_batch
+
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        if not isinstance(in_doc._backend, PdfDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(in_doc._backend).__name__} for {in_doc.file} is not a PDF backend. "
+                f"Can not convert this with a PDF pipeline. "
+                f"Please check your format configuration on DocumentConverter."
+            )
+            # conv_res.status = ConversionStatus.FAILURE
+            # return conv_res
+
+        for i in range(0, in_doc.page_count):
+            conv_res.pages.append(Page(page_no=i))
+
+        try:
+            # Iterate batches of pages (page_batch_size) in the doc
+            for page_batch in chunkify(conv_res.pages, settings.perf.page_batch_size):
+                start_pb_time = time.time()
+
+                # 1. Initialise the page resources
+                init_pages = map(
+                    functools.partial(self.initialize_page, in_doc), page_batch
+                )
+
+                # 2. Run pipeline stages
+                pipeline_pages = self._apply_on_pages(init_pages)
+
+                for p in pipeline_pages:  # Must exhaust!
+                    pass
+
+                end_pb_time = time.time() - start_pb_time
+                _log.debug(f"Finished converting page batch time={end_pb_time:.3f}")
+
+        except Exception as e:
+            conv_res.status = ConversionStatus.FAILURE
+            trace = "\n".join(traceback.format_exception(e))
+            _log.warning(
+                f"Encountered an error during conversion of document {in_doc.document_hash}:\n"
+                f"{trace}"
+            )
+            raise e
+
+        finally:
+            # Always unload the PDF backend, even in case of failure
+            if in_doc._backend:
+                in_doc._backend.unload()
+
+        return conv_res
+
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        status = ConversionStatus.SUCCESS
+        for page in conv_res.pages:
+            if page._backend is None or not page._backend.is_valid():
+                conv_res.errors.append(
+                    ErrorItem(
+                        component_type=DoclingComponentType.DOCUMENT_BACKEND,
+                        module_name=type(page._backend).__name__,
+                        error_message=f"Page {page.page_no} failed to parse.",
+                    )
+                )
+                status = ConversionStatus.PARTIAL_SUCCESS
+
+        return status
+
+    # Initialise and load resources for a page
+    @abstractmethod
+    def initialize_page(self, doc: InputDocument, page: Page) -> Page:
+        pass

docling/pipeline/simple_pipeline.py

@@ -0,0 +1,59 @@
+import logging
+
+from docling.backend.abstract_backend import (
+    AbstractDocumentBackend,
+    DeclarativeDocumentBackend,
+)
+from docling.datamodel.base_models import ConversionStatus
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.pipeline.base_pipeline import BasePipeline
+
+_log = logging.getLogger(__name__)
+
+
+class SimplePipeline(BasePipeline):
+    """SimpleModelPipeline.
+
+    This class is used at the moment for formats / backends
+    which produce straight DoclingDocument output.
+    """
+
+    def __init__(self, pipeline_options: PipelineOptions):
+        super().__init__(pipeline_options)
+
+    def _build_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+
+        if not isinstance(in_doc._backend, DeclarativeDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(in_doc._backend).__name__} for {in_doc.file} is not a declarative backend. "
+                f"Can not convert this with simple pipeline. "
+                f"Please check your format configuration on DocumentConverter."
+            )
+            # conv_res.status = ConversionStatus.FAILURE
+            # return conv_res
+
+        # Instead of running a page-level pipeline to build up the document structure,
+        # the backend is expected to be of type DeclarativeDocumentBackend, which can output
+        # a DoclingDocument straight.
+
+        conv_res.document = in_doc._backend.convert()
+        return conv_res
+
+    def _determine_status(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionStatus:
+        # This is called only if the previous steps didn't raise.
+        # Since we don't have anything else to evaluate, we can
+        # safely return SUCCESS.
+        return ConversionStatus.SUCCESS
+
+    @classmethod
+    def get_default_options(cls) -> PipelineOptions:
+        return PipelineOptions()
+
+    @classmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        return isinstance(backend, DeclarativeDocumentBackend)

docling/pipeline/standard_pdf_pipeline.py

@@ -0,0 +1,198 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from docling_core.types.doc import DocItem, ImageRef, PictureItem, TableItem
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.backend.pdf_backend import PdfDocumentBackend
+from docling.datamodel.base_models import AssembledUnit, Page
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import (
+    EasyOcrOptions,
+    PdfPipelineOptions,
+    TesseractCliOcrOptions,
+    TesseractOcrOptions,
+)
+from docling.models.base_ocr_model import BaseOcrModel
+from docling.models.ds_glm_model import GlmModel, GlmOptions
+from docling.models.easyocr_model import EasyOcrModel
+from docling.models.layout_model import LayoutModel
+from docling.models.page_assemble_model import PageAssembleModel, PageAssembleOptions
+from docling.models.page_preprocessing_model import (
+    PagePreprocessingModel,
+    PagePreprocessingOptions,
+)
+from docling.models.table_structure_model import TableStructureModel
+from docling.models.tesseract_ocr_cli_model import TesseractOcrCliModel
+from docling.models.tesseract_ocr_model import TesseractOcrModel
+from docling.pipeline.base_pipeline import PaginatedPipeline
+
+_log = logging.getLogger(__name__)
+
+
+class StandardPdfPipeline(PaginatedPipeline):
+    _layout_model_path = "model_artifacts/layout/beehive_v0.0.5_pt"
+    _table_model_path = "model_artifacts/tableformer"
+
+    def __init__(self, pipeline_options: PdfPipelineOptions):
+        super().__init__(pipeline_options)
+        self.pipeline_options: PdfPipelineOptions
+
+        if pipeline_options.artifacts_path is None:
+            self.artifacts_path = self.download_models_hf()
+        else:
+            self.artifacts_path = Path(pipeline_options.artifacts_path)
+
+        keep_images = (
+            self.pipeline_options.generate_page_images
+            or self.pipeline_options.generate_picture_images
+            or self.pipeline_options.generate_table_images
+        )
+
+        self.glm_model = GlmModel(options=GlmOptions())
+
+        if (ocr_model := self.get_ocr_model()) is None:
+            raise RuntimeError(
+                f"The specified OCR kind is not supported: {pipeline_options.ocr_options.kind}."
+            )
+
+        self.build_pipe = [
+            # Pre-processing
+            PagePreprocessingModel(
+                options=PagePreprocessingOptions(
+                    images_scale=pipeline_options.images_scale
+                )
+            ),
+            # OCR
+            ocr_model,
+            # Layout model
+            LayoutModel(
+                artifacts_path=self.artifacts_path
+                / StandardPdfPipeline._layout_model_path
+            ),
+            # Table structure model
+            TableStructureModel(
+                enabled=pipeline_options.do_table_structure,
+                artifacts_path=self.artifacts_path
+                / StandardPdfPipeline._table_model_path,
+                options=pipeline_options.table_structure_options,
+            ),
+            # Page assemble
+            PageAssembleModel(options=PageAssembleOptions(keep_images=keep_images)),
+        ]
+
+        self.enrichment_pipe = [
+            # Other models working on `NodeItem` elements in the DoclingDocument
+        ]
+
+    @staticmethod
+    def download_models_hf(
+        local_dir: Optional[Path] = None, force: bool = False
+    ) -> Path:
+        from huggingface_hub import snapshot_download
+
+        download_path = snapshot_download(
+            repo_id="ds4sd/docling-models",
+            force_download=force,
+            local_dir=local_dir,
+            revision="v2.0.1",
+        )
+
+        return Path(download_path)
+
+    def get_ocr_model(self) -> Optional[BaseOcrModel]:
+        if isinstance(self.pipeline_options.ocr_options, EasyOcrOptions):
+            return EasyOcrModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        elif isinstance(self.pipeline_options.ocr_options, TesseractCliOcrOptions):
+            return TesseractOcrCliModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        elif isinstance(self.pipeline_options.ocr_options, TesseractOcrOptions):
+            return TesseractOcrModel(
+                enabled=self.pipeline_options.do_ocr,
+                options=self.pipeline_options.ocr_options,
+            )
+        return None
+
+    def initialize_page(self, doc: InputDocument, page: Page) -> Page:
+        page._backend = doc._backend.load_page(page.page_no)  # type: ignore
+        if page._backend is not None and page._backend.is_valid():
+            page.size = page._backend.get_size()
+
+        return page
+
+    def _assemble_document(
+        self, in_doc: InputDocument, conv_res: ConversionResult
+    ) -> ConversionResult:
+        all_elements = []
+        all_headers = []
+        all_body = []
+
+        for p in conv_res.pages:
+            assert p.assembled is not None
+            for el in p.assembled.body:
+                all_body.append(el)
+            for el in p.assembled.headers:
+                all_headers.append(el)
+            for el in p.assembled.elements:
+                all_elements.append(el)
+
+        conv_res.assembled = AssembledUnit(
+            elements=all_elements, headers=all_headers, body=all_body
+        )
+
+        conv_res.document = self.glm_model(conv_res)
+
+        # Generate page images in the output
+        if self.pipeline_options.generate_page_images:
+            for page in conv_res.pages:
+                assert page.image is not None
+                page_no = page.page_no + 1
+                conv_res.document.pages[page_no].image = ImageRef.from_pil(
+                    page.image, dpi=int(72 * self.pipeline_options.images_scale)
+                )
+
+        # Generate images of the requested element types
+        if (
+            self.pipeline_options.generate_picture_images
+            or self.pipeline_options.generate_table_images
+        ):
+            scale = self.pipeline_options.images_scale
+            for element, _level in conv_res.document.iterate_items():
+                if not isinstance(element, DocItem) or len(element.prov) == 0:
+                    continue
+                if (
+                    isinstance(element, PictureItem)
+                    and self.pipeline_options.generate_picture_images
+                ) or (
+                    isinstance(element, TableItem)
+                    and self.pipeline_options.generate_table_images
+                ):
+                    page_ix = element.prov[0].page_no - 1
+                    page = conv_res.pages[page_ix]
+                    assert page.size is not None
+                    assert page.image is not None
+
+                    crop_bbox = (
+                        element.prov[0]
+                        .bbox.scaled(scale=scale)
+                        .to_top_left_origin(page_height=page.size.height * scale)
+                    )
+
+                    cropped_im = page.image.crop(crop_bbox.as_tuple())
+                    element.image = ImageRef.from_pil(cropped_im, dpi=int(72 * scale))
+
+        return conv_res
+
+    @classmethod
+    def get_default_options(cls) -> PdfPipelineOptions:
+        return PdfPipelineOptions()
+
+    @classmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        return isinstance(backend, PdfDocumentBackend)

docling/utils/export.py

@@ -1,9 +1,10 @@
 import logging
 from typing import Any, Dict, Iterable, List, Tuple, Union
 
-from docling_core.types.doc.base import BaseCell, BaseText, Ref, Table, TableCell
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.legacy_doc.base import BaseCell, BaseText, Ref, Table
 
-from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell
+from docling.datamodel.base_models import OcrCell
 from docling.datamodel.document import ConversionResult, Page
 
 _log = logging.getLogger(__name__)
@@ -40,7 +41,7 @@ def generate_multimodal_pages(
     end_ix = 0
     doc_items: List[Tuple[int, Union[BaseCell, BaseText]]] = []
 
-    doc = doc_result.output
+    doc = doc_result.legacy_document
 
     def _process_page_segments(doc_items: list[Tuple[int, BaseCell]], page: Page):
         segments = []