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

docling/models/tesseract_ocr_model.py

@@ -1,43 +1,58 @@
 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.document import ConversionResult
+from docling.datamodel.pipeline_options import TesseractOcrOptions
+from docling.datamodel.settings import settings
 from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.profiling import TimeRecorder
 
 _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
 
         if self.enabled:
-            setup_errmsg = (
+            install_errmsg = (
                 "tesserocr is not correctly installed. "
                 "Please install it via `pip install tesserocr` to use this OCR engine. "
-                "Note that tesserocr might have to be manually compiled for working with"
+                "Note that tesserocr might have to be manually compiled for working with "
                 "your Tesseract installation. The Docling documentation provides examples for it. "
-                "Alternatively, Docling has support for other OCR engines. See the documentation."
+                "Alternatively, Docling has support for other OCR engines. See the documentation: "
+                "https://ds4sd.github.io/docling/installation/"
             )
+            missing_langs_errmsg = (
+                "tesserocr is not correctly configured. No language models have been detected. "
+                "Please ensure that the TESSDATA_PREFIX envvar points to tesseract languages dir. "
+                "You can find more information how to setup other OCR engines in Docling "
+                "documentation: "
+                "https://ds4sd.github.io/docling/installation/"
+            )
+
             try:
                 import tesserocr
             except ImportError:
-                raise ImportError(setup_errmsg)
-
+                raise ImportError(install_errmsg)
             try:
                 tesseract_version = tesserocr.tesseract_version()
-                _log.debug("Initializing TesserOCR: %s", tesseract_version)
             except:
-                raise ImportError(setup_errmsg)
+                raise ImportError(install_errmsg)
+
+            _, tesserocr_languages = tesserocr.get_languages()
+            if not tesserocr_languages:
+                raise ImportError(missing_langs_errmsg)
 
             # Initialize the tesseractAPI
+            _log.debug("Initializing TesserOCR: %s", tesseract_version)
             lang = "+".join(self.options.lang)
             if self.options.path is not None:
                 self.reader = tesserocr.PyTessBaseAPI(
@@ -61,62 +76,79 @@ class TesseractOcrModel(BaseOcrModel):
             # Finalize the tesseractAPI
             self.reader.End()
 
-    def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
 
         if not self.enabled:
             yield from page_batch
             return
 
         for page in page_batch:
-            ocr_rects = self.get_ocr_rects(page)
-
-            all_ocr_cells = []
-            for ocr_rect in ocr_rects:
-                # Skip zero area boxes
-                if ocr_rect.area() == 0:
-                    continue
-                high_res_image = page._backend.get_page_image(
-                    scale=self.scale, cropbox=ocr_rect
-                )
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                yield page
+            else:
+                with TimeRecorder(conv_res, "ocr"):
 
-                # Retrieve text snippets with their bounding boxes
-                self.reader.SetImage(high_res_image)
-                boxes = self.reader.GetComponentImages(self.reader_RIL.TEXTLINE, True)
-
-                cells = []
-                for ix, (im, box, _, _) in enumerate(boxes):
-                    # Set the area of interest. Tesseract uses Bottom-Left for the origin
-                    self.reader.SetRectangle(box["x"], box["y"], box["w"], box["h"])
-
-                    # Extract text within the bounding box
-                    text = self.reader.GetUTF8Text().strip()
-                    confidence = self.reader.MeanTextConf()
-                    left = box["x"] / self.scale
-                    bottom = box["y"] / self.scale
-                    right = (box["x"] + box["w"]) / self.scale
-                    top = (box["y"] + box["h"]) / self.scale
-
-                    cells.append(
-                        OcrCell(
-                            id=ix,
-                            text=text,
-                            confidence=confidence,
-                            bbox=BoundingBox.from_tuple(
-                                coord=(left, top, right, bottom),
-                                origin=CoordOrigin.TOPLEFT,
-                            ),
+                    assert self.reader is not None
+
+                    ocr_rects = self.get_ocr_rects(page)
+
+                    all_ocr_cells = []
+                    for ocr_rect in ocr_rects:
+                        # Skip zero area boxes
+                        if ocr_rect.area() == 0:
+                            continue
+                        high_res_image = page._backend.get_page_image(
+                            scale=self.scale, cropbox=ocr_rect
                         )
-                    )
 
-                # del high_res_image
-                all_ocr_cells.extend(cells)
+                        # Retrieve text snippets with their bounding boxes
+                        self.reader.SetImage(high_res_image)
+                        boxes = self.reader.GetComponentImages(
+                            self.reader_RIL.TEXTLINE, True
+                        )
 
-            ## Remove OCR cells which overlap with programmatic cells.
-            filtered_ocr_cells = self.filter_ocr_cells(all_ocr_cells, page.cells)
+                        cells = []
+                        for ix, (im, box, _, _) in enumerate(boxes):
+                            # Set the area of interest. Tesseract uses Bottom-Left for the origin
+                            self.reader.SetRectangle(
+                                box["x"], box["y"], box["w"], box["h"]
+                            )
+
+                            # Extract text within the bounding box
+                            text = self.reader.GetUTF8Text().strip()
+                            confidence = self.reader.MeanTextConf()
+                            left = box["x"] / self.scale
+                            bottom = box["y"] / self.scale
+                            right = (box["x"] + box["w"]) / self.scale
+                            top = (box["y"] + box["h"]) / self.scale
+
+                            cells.append(
+                                OcrCell(
+                                    id=ix,
+                                    text=text,
+                                    confidence=confidence,
+                                    bbox=BoundingBox.from_tuple(
+                                        coord=(left, top, right, bottom),
+                                        origin=CoordOrigin.TOPLEFT,
+                                    ),
+                                )
+                            )
+
+                        # del high_res_image
+                        all_ocr_cells.extend(cells)
+
+                    ## Remove OCR cells which overlap with programmatic cells.
+                    filtered_ocr_cells = self.filter_ocr_cells(
+                        all_ocr_cells, page.cells
+                    )
 
-            page.cells.extend(filtered_ocr_cells)
+                    page.cells.extend(filtered_ocr_cells)
 
-            # DEBUG code:
-            # self.draw_ocr_rects_and_cells(page, ocr_rects)
+                # DEBUG code:
+                if settings.debug.visualize_ocr:
+                    self.draw_ocr_rects_and_cells(conv_res, page, ocr_rects)
 
-            yield page
+                yield page

docling/pipeline/base_pipeline.py

@@ -0,0 +1,189 @@
+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.profiling import ProfilingScope, TimeRecorder
+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:
+            with TimeRecorder(
+                conv_res, "pipeline_total", scope=ProfilingScope.DOCUMENT
+            ):
+                # These steps are building and assembling the structure of the
+                # output DoclingDocument
+                conv_res = self._build_document(conv_res)
+                conv_res = self._assemble_document(conv_res)
+                # From this stage, all operations should rely only on conv_res.output
+                conv_res = self._enrich_document(conv_res)
+                conv_res.status = self._determine_status(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, conv_res: ConversionResult) -> ConversionResult:
+        pass
+
+    def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult:
+        return conv_res
+
+    def _enrich_document(self, 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
+
+        with TimeRecorder(conv_res, "doc_enrich", scope=ProfilingScope.DOCUMENT):
+            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, 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, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        for model in self.build_pipe:
+            page_batch = model(conv_res, page_batch)
+
+        yield from page_batch
+
+    def _build_document(self, conv_res: ConversionResult) -> ConversionResult:
+
+        if not isinstance(conv_res.input._backend, PdfDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(conv_res.input._backend).__name__} for {conv_res.input.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
+
+        with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT):
+
+            for i in range(0, conv_res.input.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, conv_res), page_batch
+                    )
+
+                    # 2. Run pipeline stages
+                    pipeline_pages = self._apply_on_pages(conv_res, 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 {conv_res.input.document_hash}:\n"
+                    f"{trace}"
+                )
+                raise e
+
+            finally:
+                # Always unload the PDF backend, even in case of failure
+                if conv_res.input._backend:
+                    conv_res.input._backend.unload()
+
+        return conv_res
+
+    def _determine_status(self, 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, conv_res: ConversionResult, page: Page) -> Page:
+        pass

docling/pipeline/simple_pipeline.py

@@ -0,0 +1,56 @@
+import logging
+
+from docling.backend.abstract_backend import (
+    AbstractDocumentBackend,
+    DeclarativeDocumentBackend,
+)
+from docling.datamodel.base_models import ConversionStatus
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import PipelineOptions
+from docling.pipeline.base_pipeline import BasePipeline
+from docling.utils.profiling import ProfilingScope, TimeRecorder
+
+_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, conv_res: ConversionResult) -> ConversionResult:
+
+        if not isinstance(conv_res.input._backend, DeclarativeDocumentBackend):
+            raise RuntimeError(
+                f"The selected backend {type(conv_res.input._backend).__name__} for {conv_res.input.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.
+        with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT):
+            conv_res.document = conv_res.input._backend.convert()
+        return conv_res
+
+    def _determine_status(self, 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,201 @@
+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
+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
+from docling.utils.profiling import ProfilingScope, TimeRecorder
+
+_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, conv_res: ConversionResult, page: Page) -> Page:
+        with TimeRecorder(conv_res, "page_init"):
+            page._backend = conv_res.input._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, conv_res: ConversionResult) -> ConversionResult:
+        all_elements = []
+        all_headers = []
+        all_body = []
+
+        with TimeRecorder(conv_res, "doc_assemble", scope=ProfilingScope.DOCUMENT):
+            for p in conv_res.pages:
+                if 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 = []