docling 2.69.0__py3-none-any.whl

docling/models/stages/ocr/tesseract_ocr_model.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Iterable, Optional, Type
+
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.doc.page import TextCell
+
+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,
+    TesseractOcrOptions,
+)
+from docling.datamodel.settings import settings
+from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.ocr_utils import (
+    map_tesseract_script,
+    parse_tesseract_orientation,
+    tesseract_box_to_bounding_rectangle,
+)
+from docling.utils.profiling import TimeRecorder
+
+_log = logging.getLogger(__name__)
+
+
+class TesseractOcrModel(BaseOcrModel):
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        options: TesseractOcrOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        super().__init__(
+            enabled=enabled,
+            artifacts_path=artifacts_path,
+            options=options,
+            accelerator_options=accelerator_options,
+        )
+        self.options: TesseractOcrOptions
+        self._is_auto: bool = "auto" in self.options.lang
+        self.scale = 3  # multiplier for 72 dpi == 216 dpi.
+        self.reader = None
+        self.script_readers: dict[str, tesserocr.PyTessBaseAPI] = {}
+
+        if self.enabled:
+            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 "
+                "your Tesseract installation. The Docling documentation provides examples for it. "
+                "Alternatively, Docling has support for other OCR engines. See the documentation: "
+                "https://docling-project.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://docling-project.github.io/docling/installation/"
+            )
+
+            try:
+                import tesserocr
+            except ImportError:
+                raise ImportError(install_errmsg)
+            try:
+                tesseract_version = tesserocr.tesseract_version()
+            except Exception:
+                raise ImportError(install_errmsg)
+
+            _, self._tesserocr_languages = tesserocr.get_languages()
+            if not self._tesserocr_languages:
+                raise ImportError(missing_langs_errmsg)
+
+            # Initialize the tesseractAPI
+            _log.debug("Initializing TesserOCR: %s", tesseract_version)
+            lang = "+".join(self.options.lang)
+
+            if any(lang.startswith("script/") for lang in self._tesserocr_languages):
+                self.script_prefix = "script/"
+            else:
+                self.script_prefix = ""
+
+            tesserocr_kwargs = {
+                "init": True,
+                "oem": tesserocr.OEM.DEFAULT,
+            }
+
+            self.osd_reader = None
+
+            if self.options.path is not None:
+                tesserocr_kwargs["path"] = self.options.path
+
+            # Set main OCR reader with configurable PSM
+            main_psm = (
+                self.options.psm if self.options.psm is not None else tesserocr.PSM.AUTO
+            )
+            if lang == "auto":
+                self.reader = tesserocr.PyTessBaseAPI(psm=main_psm, **tesserocr_kwargs)
+            else:
+                self.reader = tesserocr.PyTessBaseAPI(
+                    lang=lang,
+                    psm=main_psm,
+                    **tesserocr_kwargs,
+                )
+            # OSD reader must use PSM.OSD_ONLY for orientation detection
+            self.osd_reader = tesserocr.PyTessBaseAPI(
+                lang="osd", psm=tesserocr.PSM.OSD_ONLY, **tesserocr_kwargs
+            )
+            self.reader_RIL = tesserocr.RIL
+
+    def __del__(self):
+        if self.reader is not None:
+            # Finalize the tesseractAPI
+            self.reader.End()
+        for script in self.script_readers:
+            self.script_readers[script].End()
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        if not self.enabled:
+            yield from page_batch
+            return
+
+        for page_i, page in enumerate(page_batch):
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                yield page
+            else:
+                with TimeRecorder(conv_res, "ocr"):
+                    assert self.reader is not None
+                    assert self.osd_reader is not None
+                    assert self._tesserocr_languages is not None
+
+                    ocr_rects = self.get_ocr_rects(page)
+
+                    all_ocr_cells = []
+                    for ocr_rect_i, ocr_rect in enumerate(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
+                        )
+
+                        local_reader = self.reader
+                        self.osd_reader.SetImage(high_res_image)
+
+                        doc_orientation = 0
+                        osd = self.osd_reader.DetectOrientationScript()
+
+                        # No text, or Orientation and Script detection failure
+                        if osd is None:
+                            _log.error(
+                                "OSD failed for doc (doc %s, page: %s, "
+                                "OCR rectangle: %s)",
+                                conv_res.input.file,
+                                page_i,
+                                ocr_rect_i,
+                            )
+                            # Skipping if OSD fail when in auto mode, otherwise proceed
+                            # to OCR in the hope OCR will succeed while OSD failed
+                            if self._is_auto:
+                                continue
+                        else:
+                            doc_orientation = parse_tesseract_orientation(
+                                osd["orient_deg"]
+                            )
+                            if doc_orientation != 0:
+                                high_res_image = high_res_image.rotate(
+                                    -doc_orientation, expand=True
+                                )
+                        if self._is_auto:
+                            script = osd["script_name"]
+                            script = map_tesseract_script(script)
+                            lang = f"{self.script_prefix}{script}"
+
+                            # Check if the detected language is present in the system
+                            if lang not in self._tesserocr_languages:
+                                msg = f"Tesseract detected the script '{script}' and language '{lang}'."
+                                msg += " However this language is not installed in your system and will be ignored."
+                                _log.warning(msg)
+                            else:
+                                if script not in self.script_readers:
+                                    import tesserocr
+
+                                    self.script_readers[script] = (
+                                        tesserocr.PyTessBaseAPI(
+                                            path=self.reader.GetDatapath(),
+                                            lang=lang,
+                                            psm=self.options.psm
+                                            if self.options.psm is not None
+                                            else tesserocr.PSM.AUTO,
+                                            init=True,
+                                            oem=tesserocr.OEM.DEFAULT,
+                                        )
+                                    )
+                                local_reader = self.script_readers[script]
+
+                        local_reader.SetImage(high_res_image)
+                        boxes = local_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
+                            local_reader.SetRectangle(
+                                box["x"], box["y"], box["w"], box["h"]
+                            )
+
+                            # Extract text within the bounding box
+                            text = local_reader.GetUTF8Text().strip()
+                            confidence = local_reader.MeanTextConf()
+                            left, top = box["x"], box["y"]
+                            right = left + box["w"]
+                            bottom = top + box["h"]
+                            bbox = BoundingBox(
+                                l=left,
+                                t=top,
+                                r=right,
+                                b=bottom,
+                                coord_origin=CoordOrigin.TOPLEFT,
+                            )
+                            rect = tesseract_box_to_bounding_rectangle(
+                                bbox,
+                                original_offset=ocr_rect,
+                                scale=self.scale,
+                                orientation=doc_orientation,
+                                im_size=high_res_image.size,
+                            )
+                            cells.append(
+                                TextCell(
+                                    index=ix,
+                                    text=text,
+                                    orig=text,
+                                    from_ocr=True,
+                                    confidence=confidence,
+                                    rect=rect,
+                                )
+                            )
+
+                        # del high_res_image
+                        all_ocr_cells.extend(cells)
+
+                    # Post-process the cells
+                    self.post_process_cells(all_ocr_cells, page)
+
+                # DEBUG code:
+                if settings.debug.visualize_ocr:
+                    self.draw_ocr_rects_and_cells(conv_res, page, ocr_rects)
+
+                yield page
+
+    @classmethod
+    def get_options_type(cls) -> Type[OcrOptions]:
+        return TesseractOcrOptions

docling/models/stages/page_assemble/page_assemble_model.py

@@ -0,0 +1,156 @@
+import logging
+import re
+from collections.abc import Iterable
+from typing import List
+
+import numpy as np
+from pydantic import BaseModel
+
+from docling.datamodel.base_models import (
+    AssembledUnit,
+    ContainerElement,
+    FigureElement,
+    Page,
+    PageElement,
+    Table,
+    TextElement,
+)
+from docling.datamodel.document import ConversionResult
+from docling.models.base_model import BasePageModel
+from docling.models.stages.layout.layout_model import LayoutModel
+from docling.utils.profiling import TimeRecorder
+
+_log = logging.getLogger(__name__)
+
+
+class PageAssembleOptions(BaseModel):
+    pass
+
+
+class PageAssembleModel(BasePageModel):
+    def __init__(self, options: PageAssembleOptions):
+        self.options = options
+
+    def sanitize_text(self, lines):
+        if len(lines) <= 1:
+            return " ".join(lines)
+
+        for ix, line in enumerate(lines[1:]):
+            prev_line = lines[ix]
+
+            if prev_line.endswith("-"):
+                prev_words = re.findall(r"\b[\w]+\b", prev_line)
+                line_words = re.findall(r"\b[\w]+\b", line)
+
+                if (
+                    len(prev_words)
+                    and len(line_words)
+                    and prev_words[-1].isalnum()
+                    and line_words[0].isalnum()
+                ):
+                    lines[ix] = prev_line[:-1]
+            else:
+                lines[ix] += " "
+
+        sanitized_text = "".join(lines)
+
+        # Text normalization
+        sanitized_text = sanitized_text.replace("⁄", "/")  # noqa: RUF001
+        sanitized_text = sanitized_text.replace("’", "'")  # noqa: RUF001
+        sanitized_text = sanitized_text.replace("‘", "'")  # noqa: RUF001
+        sanitized_text = sanitized_text.replace("“", '"')
+        sanitized_text = sanitized_text.replace("”", '"')
+        sanitized_text = sanitized_text.replace("•", "·")
+
+        return sanitized_text.strip()  # Strip any leading or trailing whitespace
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        for page in page_batch:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                yield page
+            else:
+                with TimeRecorder(conv_res, "page_assemble"):
+                    assert page.predictions.layout is not None
+
+                    # assembles some JSON output page by page.
+
+                    elements: List[PageElement] = []
+                    headers: List[PageElement] = []
+                    body: List[PageElement] = []
+
+                    for cluster in page.predictions.layout.clusters:
+                        # _log.info("Cluster label seen:", cluster.label)
+                        if cluster.label in LayoutModel.TEXT_ELEM_LABELS:
+                            textlines = [
+                                cell.text.replace("\x02", "-").strip()
+                                for cell in cluster.cells
+                                if len(cell.text.strip()) > 0
+                            ]
+                            text = self.sanitize_text(textlines)
+                            text_el = TextElement(
+                                label=cluster.label,
+                                id=cluster.id,
+                                text=text,
+                                page_no=page.page_no,
+                                cluster=cluster,
+                            )
+                            elements.append(text_el)
+
+                            if cluster.label in LayoutModel.PAGE_HEADER_LABELS:
+                                headers.append(text_el)
+                            else:
+                                body.append(text_el)
+                        elif cluster.label in LayoutModel.TABLE_LABELS:
+                            tbl = None
+                            if page.predictions.tablestructure:
+                                tbl = page.predictions.tablestructure.table_map.get(
+                                    cluster.id, None
+                                )
+                            if not tbl:  # fallback: add table without structure, if it isn't present
+                                tbl = Table(
+                                    label=cluster.label,
+                                    id=cluster.id,
+                                    text="",
+                                    otsl_seq=[],
+                                    table_cells=[],
+                                    cluster=cluster,
+                                    page_no=page.page_no,
+                                )
+
+                            elements.append(tbl)
+                            body.append(tbl)
+                        elif cluster.label == LayoutModel.FIGURE_LABEL:
+                            fig = None
+                            if page.predictions.figures_classification:
+                                fig = page.predictions.figures_classification.figure_map.get(
+                                    cluster.id, None
+                                )
+                            if not fig:  # fallback: add figure without classification, if it isn't present
+                                fig = FigureElement(
+                                    label=cluster.label,
+                                    id=cluster.id,
+                                    text="",
+                                    data=None,
+                                    cluster=cluster,
+                                    page_no=page.page_no,
+                                )
+                            elements.append(fig)
+                            body.append(fig)
+                        elif cluster.label in LayoutModel.CONTAINER_LABELS:
+                            container_el = ContainerElement(
+                                label=cluster.label,
+                                id=cluster.id,
+                                page_no=page.page_no,
+                                cluster=cluster,
+                            )
+                            elements.append(container_el)
+                            body.append(container_el)
+
+                    page.assembled = AssembledUnit(
+                        elements=elements, headers=headers, body=body
+                    )
+
+                yield page

docling/models/stages/page_preprocessing/page_preprocessing_model.py

@@ -0,0 +1,145 @@
+import re
+import warnings
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Literal, Optional
+
+import numpy as np
+from PIL import ImageDraw
+from pydantic import BaseModel
+
+from docling.datamodel.base_models import Page
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.settings import settings
+from docling.models.base_model import BasePageModel
+from docling.utils.profiling import TimeRecorder
+
+
+class PagePreprocessingOptions(BaseModel):
+    images_scale: Optional[float]
+    skip_cell_extraction: bool = (
+        False  # Skip text cell extraction for VLM-only processing
+    )
+
+
+class PagePreprocessingModel(BasePageModel):
+    def __init__(self, options: PagePreprocessingOptions):
+        self.options = options
+
+        # Pre-compiled regex patterns for efficiency
+        self.GLYPH_RE = re.compile(r"GLYPH<[0-9A-Fa-f]+>")
+        self.SLASH_G_RE = re.compile(r"(?:/G\d+){2,}")
+        self.FRAG_RE = re.compile(r"\b[A-Za-z](?:/[a-z]{1,3}\.[a-z]{1,3}){2,}\b")
+        self.SLASH_NUMBER_GARBAGE_RE = re.compile(
+            r"(?:/\w+\s*){2,}"
+        )  # Two or more "/token " sequences
+
+    def __call__(
+        self, conv_res: ConversionResult, page_batch: Iterable[Page]
+    ) -> Iterable[Page]:
+        for page in page_batch:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                yield page
+            else:
+                with TimeRecorder(conv_res, "page_parse"):
+                    page = self._populate_page_images(page)
+                    if not self.options.skip_cell_extraction:
+                        page = self._parse_page_cells(conv_res, page)
+                yield page
+
+    # Generate the page image and store it in the page object
+    def _populate_page_images(self, page: Page) -> Page:
+        # default scale
+        page.get_image(
+            scale=1.0
+        )  # puts the page image on the image cache at default scale
+
+        images_scale = self.options.images_scale
+        # user requested scales
+        if images_scale is not None:
+            page._default_image_scale = images_scale
+            page.get_image(
+                scale=images_scale
+            )  # this will trigger storing the image in the internal cache
+
+        return page
+
+    # Extract and populate the page cells and store it in the page object
+    def _parse_page_cells(self, conv_res: ConversionResult, page: Page) -> Page:
+        assert page._backend is not None
+
+        page.parsed_page = page._backend.get_segmented_page()
+        assert page.parsed_page is not None
+
+        # Rate the text quality from the PDF parser, and aggregate on page
+        text_scores = []
+        for c in page.cells:
+            score = self.rate_text_quality(c.text)
+            text_scores.append(score)
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", "Mean of empty slice", RuntimeWarning, "numpy"
+            )
+            conv_res.confidence.pages[page.page_no].parse_score = float(
+                np.nanquantile(
+                    text_scores, q=0.10
+                )  # To emphasise problems in the parse_score, we take the 10% percentile score of all text cells.
+            )
+
+        # DEBUG code:
+        def draw_text_boxes(image, cells, show: bool = False):
+            draw = ImageDraw.Draw(image.copy())
+            for c in cells:
+                x0, y0, x1, y1 = (
+                    c.to_bounding_box().l,
+                    c.to_bounding_box().t,
+                    c.to_bounding_box().r,
+                    c.to_bounding_box().b,
+                )
+
+                draw.rectangle([(x0, y0), (x1, y1)], outline="red")
+            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"cells_page_{page.page_no:05}.png"
+                image.save(str(out_file), format="png")
+
+        if settings.debug.visualize_cells:
+            draw_text_boxes(page.get_image(scale=1.0), page.cells)
+
+        return page
+
+    def rate_text_quality(self, text: str) -> float:
+        # Hard errors: if any of these patterns are found, return 0.0 immediately.
+        blacklist_chars = ["�"]
+        if (
+            any(text.find(c) >= 0 for c in blacklist_chars)
+            or self.GLYPH_RE.search(text)
+            or self.SLASH_G_RE.search(text)
+            or self.SLASH_NUMBER_GARBAGE_RE.match(
+                text
+            )  # Check if text is mostly slash-number pattern
+        ):
+            return 0.0
+
+        penalty = 0.0
+
+        # Apply a penalty only if the fragmented words pattern occurs at least three times.
+        frag_matches = self.FRAG_RE.findall(text)
+        if len(frag_matches) >= 3:
+            penalty += 0.1 * len(frag_matches)
+
+        # Additional heuristic: if the average token length is below 2, add a penalty.
+        # tokens = text.split()
+        # if tokens and (sum(map(len, tokens)) / len(tokens)) < 2:
+        #    penalty += 0.2
+
+        return max(1.0 - penalty, 0.0)