docling 2.69.0__py3-none-any.whl

docling/models/base_model.py

@@ -0,0 +1,230 @@
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Any, Generic, Optional, Protocol, Type, Union
+
+import numpy as np
+from docling_core.types.doc import (
+    BoundingBox,
+    DocItem,
+    DoclingDocument,
+    NodeItem,
+    PictureItem,
+)
+from PIL.Image import Image
+from typing_extensions import TypeVar
+
+from docling.datamodel.base_models import (
+    ItemAndImageEnrichmentElement,
+    Page,
+    VlmPrediction,
+)
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import BaseOptions
+from docling.datamodel.pipeline_options_vlm_model import (
+    InlineVlmOptions,
+    TransformersPromptStyle,
+)
+from docling.datamodel.settings import settings
+
+
+class BaseModelWithOptions(Protocol):
+    @classmethod
+    def get_options_type(cls) -> Type[BaseOptions]: ...
+
+    def __init__(self, *, options: BaseOptions, **kwargs): ...
+
+
+class BasePageModel(ABC):
+    @abstractmethod
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        pass
+
+
+class BaseVlmModel(ABC):
+    """Base class for Vision-Language Models that adds image processing capability."""
+
+    @abstractmethod
+    def process_images(
+        self,
+        image_batch: Iterable[Union[Image, np.ndarray]],
+        prompt: Union[str, list[str]],
+    ) -> Iterable[VlmPrediction]:
+        """Process raw images without page metadata.
+
+        Args:
+            image_batch: Iterable of PIL Images or numpy arrays
+            prompt: Either:
+                - str: Single prompt used for all images
+                - list[str]: List of prompts (one per image, must match image count)
+
+        Raises:
+            ValueError: If prompt list length doesn't match image count.
+        """
+
+
+class BaseVlmPageModel(BasePageModel, BaseVlmModel):
+    """Base implementation for VLM models that inherit from BasePageModel.
+
+    Provides a default __call__ implementation that extracts images from pages,
+    processes them using process_images, and attaches results back to pages.
+    """
+
+    # Type annotations for attributes that subclasses must initialize
+    vlm_options: InlineVlmOptions
+    processor: Any
+
+    def _build_prompt_safe(self, page: Page) -> str:
+        """Build prompt with backward compatibility for user overrides.
+
+        Tries to call build_prompt with _internal_page parameter (for layout-aware
+        pipelines). Falls back to basic call if user override doesn't accept it.
+
+        Args:
+            page: The full Page object with layout predictions and parsed_page.
+
+        Returns:
+            The formatted prompt string.
+        """
+        try:
+            return self.vlm_options.build_prompt(page.parsed_page, _internal_page=page)
+        except TypeError:
+            # User override doesn't accept _internal_page - fall back to basic call
+            return self.vlm_options.build_prompt(page.parsed_page)
+
+    @abstractmethod
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        """Extract images from pages, process them, and attach results back."""
+
+    def formulate_prompt(self, user_prompt: str) -> str:
+        """Formulate a prompt for the VLM."""
+        _log = logging.getLogger(__name__)
+
+        if self.vlm_options.transformers_prompt_style == TransformersPromptStyle.RAW:
+            return user_prompt
+        elif self.vlm_options.transformers_prompt_style == TransformersPromptStyle.NONE:
+            return ""
+        elif self.vlm_options.repo_id == "microsoft/Phi-4-multimodal-instruct":
+            _log.debug("Using specialized prompt for Phi-4")
+            # Note: This might need adjustment for VLLM vs transformers
+            user_prompt_prefix = "<|user|>"
+            assistant_prompt = "<|assistant|>"
+            prompt_suffix = "<|end|>"
+
+            prompt = f"{user_prompt_prefix}<|image_1|>{user_prompt}{prompt_suffix}{assistant_prompt}"
+            _log.debug(f"prompt for {self.vlm_options.repo_id}: {prompt}")
+
+            return prompt
+
+        elif self.vlm_options.transformers_prompt_style == TransformersPromptStyle.CHAT:
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "This is a page from a document.",
+                        },
+                        {"type": "image"},
+                        {"type": "text", "text": user_prompt},
+                    ],
+                }
+            ]
+            prompt = self.processor.apply_chat_template(
+                messages, add_generation_prompt=True
+            )
+            return prompt
+
+        raise RuntimeError(
+            f"Unknown prompt style `{self.vlm_options.transformers_prompt_style}`. Valid values are {', '.join(s.value for s in TransformersPromptStyle)}."
+        )
+
+
+EnrichElementT = TypeVar("EnrichElementT", default=NodeItem)
+
+
+class GenericEnrichmentModel(ABC, Generic[EnrichElementT]):
+    elements_batch_size: int = settings.perf.elements_batch_size
+
+    @abstractmethod
+    def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
+        pass
+
+    @abstractmethod
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[EnrichElementT]:
+        pass
+
+    @abstractmethod
+    def __call__(
+        self, doc: DoclingDocument, element_batch: Iterable[EnrichElementT]
+    ) -> Iterable[NodeItem]:
+        pass
+
+
+class BaseEnrichmentModel(GenericEnrichmentModel[NodeItem]):
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[NodeItem]:
+        if self.is_processable(doc=conv_res.document, element=element):
+            return element
+        return None
+
+
+class BaseItemAndImageEnrichmentModel(
+    GenericEnrichmentModel[ItemAndImageEnrichmentElement]
+):
+    images_scale: float
+    expansion_factor: float = 0.0
+
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[ItemAndImageEnrichmentElement]:
+        if not self.is_processable(doc=conv_res.document, element=element):
+            return None
+
+        assert isinstance(element, DocItem)
+
+        # Allow the case of documents without page images but embedded images (e.g. Word and HTML docs)
+        if isinstance(element, PictureItem):
+            embedded_im = element.get_image(conv_res.document)
+            if embedded_im is not None:
+                return ItemAndImageEnrichmentElement(item=element, image=embedded_im)
+            elif len(element.prov) == 0:
+                return None
+
+        # Crop the image form the page
+        element_prov = element.prov[0]
+        bbox = element_prov.bbox
+        width = bbox.r - bbox.l
+        height = bbox.t - bbox.b
+
+        # TODO: move to a utility in the BoundingBox class
+        expanded_bbox = BoundingBox(
+            l=bbox.l - width * self.expansion_factor,
+            t=bbox.t + height * self.expansion_factor,
+            r=bbox.r + width * self.expansion_factor,
+            b=bbox.b - height * self.expansion_factor,
+            coord_origin=bbox.coord_origin,
+        )
+
+        page_ix = element_prov.page_no - conv_res.pages[0].page_no
+        cropped_image = conv_res.pages[page_ix].get_image(
+            scale=self.images_scale, cropbox=expanded_bbox
+        )
+
+        # Allow for images being embedded without the page backend or page images
+        if cropped_image is None and isinstance(element, PictureItem):
+            embedded_im = element.get_image(conv_res.document)
+            if embedded_im is not None:
+                return ItemAndImageEnrichmentElement(item=element, image=embedded_im)
+            else:
+                return None
+
+        # Return the proper cropped image
+        return ItemAndImageEnrichmentElement(item=element, image=cropped_image)

docling/models/base_ocr_model.py

@@ -0,0 +1,241 @@
+import copy
+import logging
+from abc import abstractmethod
+from collections.abc import Iterable
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Type
+
+import numpy as np
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.doc.page import TextCell
+from PIL import Image, ImageDraw
+from rtree import index
+
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.base_models import Page
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import OcrOptions
+from docling.datamodel.settings import settings
+from docling.models.base_model import BaseModelWithOptions, BasePageModel
+
+_log = logging.getLogger(__name__)
+
+
+class BaseOcrModel(BasePageModel, BaseModelWithOptions):
+    def __init__(
+        self,
+        *,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        options: OcrOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        # Make sure any delay/error from import occurs on ocr model init and not first use
+        from scipy.ndimage import binary_dilation, find_objects, label
+
+        self.enabled = enabled
+        self.options = options
+
+    # Computes the optimum amount and coordinates of rectangles to OCR on a given page
+    def get_ocr_rects(self, page: Page) -> List[BoundingBox]:
+        from scipy.ndimage import binary_dilation, find_objects, label
+
+        BITMAP_COVERAGE_TRESHOLD = 0.75
+        assert page.size is not None
+
+        def find_ocr_rects(size, bitmap_rects):
+            image = Image.new(
+                "1", (round(size.width), round(size.height))
+            )  # '1' mode is binary
+
+            # Draw all bitmap rects into a binary image
+            draw = ImageDraw.Draw(image)
+            for rect in bitmap_rects:
+                x0, y0, x1, y1 = rect.as_tuple()
+                x0, y0, x1, y1 = round(x0), round(y0), round(x1), round(y1)
+                draw.rectangle([(x0, y0), (x1, y1)], fill=1)
+
+            np_image = np.array(image)
+
+            # Dilate the image by 10 pixels to merge nearby bitmap rectangles
+            structure = np.ones(
+                (20, 20)
+            )  # Create a 20x20 structure element (10 pixels in all directions)
+            np_image = binary_dilation(np_image > 0, structure=structure)
+
+            # Find the connected components
+            labeled_image, num_features = label(
+                np_image > 0
+            )  # Label black (0 value) regions
+
+            # Find enclosing bounding boxes for each connected component.
+            slices = find_objects(labeled_image)
+            bounding_boxes = [
+                BoundingBox(
+                    l=slc[1].start,
+                    t=slc[0].start,
+                    r=slc[1].stop - 1,
+                    b=slc[0].stop - 1,
+                    coord_origin=CoordOrigin.TOPLEFT,
+                )
+                for slc in slices
+            ]
+
+            # Compute area fraction on page covered by bitmaps
+            area_frac = np.sum(np_image > 0) / (size.width * size.height)
+
+            return (area_frac, bounding_boxes)  # fraction covered  # boxes
+
+        if page._backend is not None:
+            bitmap_rects = page._backend.get_bitmap_rects()
+        else:
+            bitmap_rects = []
+        coverage, ocr_rects = find_ocr_rects(page.size, bitmap_rects)
+
+        # return full-page rectangle if page is dominantly covered with bitmaps
+        if self.options.force_full_page_ocr or coverage > max(
+            BITMAP_COVERAGE_TRESHOLD, self.options.bitmap_area_threshold
+        ):
+            return [
+                BoundingBox(
+                    l=0,
+                    t=0,
+                    r=page.size.width,
+                    b=page.size.height,
+                    coord_origin=CoordOrigin.TOPLEFT,
+                )
+            ]
+        # return individual rectangles if the bitmap coverage is above the threshold
+        elif coverage > self.options.bitmap_area_threshold:
+            return ocr_rects
+        else:  # overall coverage of bitmaps is too low, drop all bitmap rectangles.
+            return []
+
+    # Filters OCR cells by dropping any OCR cell that intersects with an existing programmatic cell.
+    def _filter_ocr_cells(
+        self, ocr_cells: List[TextCell], programmatic_cells: List[TextCell]
+    ) -> List[TextCell]:
+        # Create R-tree index for programmatic cells
+        p = index.Property()
+        p.dimension = 2
+        idx = index.Index(properties=p)
+        for i, cell in enumerate(programmatic_cells):
+            idx.insert(i, cell.rect.to_bounding_box().as_tuple())
+
+        def is_overlapping_with_existing_cells(ocr_cell):
+            # Query the R-tree to get overlapping rectangles
+            possible_matches_index = list(
+                idx.intersection(ocr_cell.rect.to_bounding_box().as_tuple())
+            )
+
+            return (
+                len(possible_matches_index) > 0
+            )  # this is a weak criterion but it works.
+
+        filtered_ocr_cells = [
+            rect for rect in ocr_cells if not is_overlapping_with_existing_cells(rect)
+        ]
+        return filtered_ocr_cells
+
+    def post_process_cells(self, ocr_cells: List[TextCell], page: Page) -> None:
+        r"""
+        Post-process the OCR cells and update the page object.
+        Updates parsed_page.textline_cells directly since page.cells is now read-only.
+        """
+        # Get existing cells from the read-only property
+        existing_cells = page.cells
+
+        # Combine existing and OCR cells with overlap filtering
+        final_cells = self._combine_cells(existing_cells, ocr_cells)
+
+        assert page.parsed_page is not None
+
+        # Update parsed_page.textline_cells directly
+        page.parsed_page.textline_cells = final_cells
+        page.parsed_page.has_lines = len(final_cells) > 0
+
+        # When force_full_page_ocr is used, PDF-extracted word/char cells are
+        # unreliable. Filter out cells where from_ocr=False, keeping any OCR-
+        # generated cells. This ensures downstream components (e.g., table
+        # structure model) fall back to OCR-extracted textline cells.
+        if self.options.force_full_page_ocr:
+            page.parsed_page.word_cells = [
+                c for c in page.parsed_page.word_cells if c.from_ocr
+            ]
+            page.parsed_page.char_cells = [
+                c for c in page.parsed_page.char_cells if c.from_ocr
+            ]
+            page.parsed_page.has_words = len(page.parsed_page.word_cells) > 0
+            page.parsed_page.has_chars = len(page.parsed_page.char_cells) > 0
+
+    def _combine_cells(
+        self, existing_cells: List[TextCell], ocr_cells: List[TextCell]
+    ) -> List[TextCell]:
+        """Combine existing and OCR cells with filtering and re-indexing."""
+        if self.options.force_full_page_ocr:
+            combined = ocr_cells
+        else:
+            filtered_ocr_cells = self._filter_ocr_cells(ocr_cells, existing_cells)
+            combined = list(existing_cells) + filtered_ocr_cells
+
+        # Re-index in-place
+        for i, cell in enumerate(combined):
+            cell.index = i
+
+        return combined
+
+    def draw_ocr_rects_and_cells(self, conv_res, page, ocr_rects, show: bool = False):
+        image = copy.deepcopy(page.image)
+        scale_x = image.width / page.size.width
+        scale_y = image.height / page.size.height
+
+        draw = ImageDraw.Draw(image, "RGBA")
+
+        # Draw OCR rectangles as yellow filled rect
+        for rect in ocr_rects:
+            x0, y0, x1, y1 = rect.as_tuple()
+            y0 *= scale_x
+            y1 *= scale_y
+            x0 *= scale_x
+            x1 *= scale_x
+
+            shade_color = (255, 255, 0, 40)  # transparent yellow
+            draw.rectangle([(x0, y0), (x1, y1)], fill=shade_color, outline=None)
+
+        # Draw OCR and programmatic cells
+        for tc in page.cells:
+            x0, y0, x1, y1 = tc.rect.to_bounding_box().as_tuple()
+            y0 *= scale_x
+            y1 *= scale_y
+            x0 *= scale_x
+            x1 *= scale_x
+
+            if y1 <= y0:
+                y1, y0 = y0, y1
+
+            color = "magenta" if tc.from_ocr else "gray"
+
+            draw.rectangle([(x0, y0), (x1, y1)], outline=color)
+
+        if show:
+            image.show()
+        else:
+            out_path: Path = (
+                Path(settings.debug.debug_output_path)
+                / f"debug_{conv_res.input.file.stem}"
+            )
+            out_path.mkdir(parents=True, exist_ok=True)
+
+            out_file = out_path / f"ocr_page_{page.page_no:05}.png"
+            image.save(str(out_file), format="png")
+
+    @abstractmethod
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_options_type(cls) -> Type[OcrOptions]:
+        pass

docling/models/base_table_model.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from typing import Type
+
+from docling.datamodel.base_models import Page, TableStructurePrediction
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import BaseTableStructureOptions
+from docling.models.base_model import BaseModelWithOptions, BasePageModel
+
+
+class BaseTableStructureModel(BasePageModel, BaseModelWithOptions, ABC):
+    """Shared interface for table structure models."""
+
+    enabled: bool
+
+    @classmethod
+    @abstractmethod
+    def get_options_type(cls) -> Type[BaseTableStructureOptions]:
+        """Return the options type supported by this table model."""
+
+    @abstractmethod
+    def predict_tables(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[TableStructurePrediction]:
+        """Produce table structure predictions for the provided pages."""
+
+    def __call__(
+        self,
+        conv_res: ConversionResult,
+        page_batch: Iterable[Page],
+    ) -> Iterable[Page]:
+        if not getattr(self, "enabled", True):
+            yield from page_batch
+            return
+
+        pages = list(page_batch)
+        predictions = self.predict_tables(conv_res, pages)
+
+        for page, prediction in zip(pages, predictions):
+            page.predictions.tablestructure = prediction
+            yield page