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

docling/models/easyocr_model.py

@@ -5,8 +5,11 @@ import numpy
 from docling_core.types.doc import BoundingBox, CoordOrigin
 
 from docling.datamodel.base_models import OcrCell, Page
+from docling.datamodel.document import ConversionResult
 from docling.datamodel.pipeline_options import EasyOcrOptions
+from docling.datamodel.settings import settings
 from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
 
@@ -33,58 +36,65 @@ class EasyOcrModel(BaseOcrModel):
                 download_enabled=self.options.download_enabled,
             )
 
-    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:
+
             assert page._backend is not None
             if not page._backend.is_valid():
                 yield page
             else:
-                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
-                    )
-                    im = numpy.array(high_res_image)
-                    result = self.reader.readtext(im)
-
-                    del high_res_image
-                    del im
-
-                    cells = [
-                        OcrCell(
-                            id=ix,
-                            text=line[1],
-                            confidence=line[2],
-                            bbox=BoundingBox.from_tuple(
-                                coord=(
-                                    (line[0][0][0] / self.scale) + ocr_rect.l,
-                                    (line[0][0][1] / self.scale) + ocr_rect.t,
-                                    (line[0][2][0] / self.scale) + ocr_rect.l,
-                                    (line[0][2][1] / self.scale) + ocr_rect.t,
-                                ),
-                                origin=CoordOrigin.TOPLEFT,
-                            ),
+                with TimeRecorder(conv_res, "ocr"):
+                    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
                         )
-                        for ix, line in enumerate(result)
-                    ]
-                    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)
+                        im = numpy.array(high_res_image)
+                        result = self.reader.readtext(im)
+
+                        del high_res_image
+                        del im
+
+                        cells = [
+                            OcrCell(
+                                id=ix,
+                                text=line[1],
+                                confidence=line[2],
+                                bbox=BoundingBox.from_tuple(
+                                    coord=(
+                                        (line[0][0][0] / self.scale) + ocr_rect.l,
+                                        (line[0][0][1] / self.scale) + ocr_rect.t,
+                                        (line[0][2][0] / self.scale) + ocr_rect.l,
+                                        (line[0][2][1] / self.scale) + ocr_rect.t,
+                                    ),
+                                    origin=CoordOrigin.TOPLEFT,
+                                ),
+                            )
+                            for ix, line in enumerate(result)
+                        ]
+                        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)
+                if settings.debug.visualize_ocr:
+                    self.draw_ocr_rects_and_cells(conv_res, page, ocr_rects)
 
                 yield page

docling/models/layout_model.py

@@ -16,8 +16,11 @@ from docling.datamodel.base_models import (
     LayoutPrediction,
     Page,
 )
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.settings import settings
 from docling.models.base_model import BasePageModel
 from docling.utils import layout_utils as lu
+from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
 
@@ -271,74 +274,97 @@ class LayoutModel(BasePageModel):
 
         return clusters_out_new, cells_out_new
 
-    def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+    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:
-                assert page.size is not None
-
-                clusters = []
-                for ix, pred_item in enumerate(
-                    self.layout_predictor.predict(page.get_image(scale=1.0))
-                ):
-                    label = DocItemLabel(
-                        pred_item["label"].lower().replace(" ", "_").replace("-", "_")
-                    )  # Temporary, until docling-ibm-model uses docling-core types
-                    cluster = Cluster(
-                        id=ix,
-                        label=label,
-                        confidence=pred_item["confidence"],
-                        bbox=BoundingBox.model_validate(pred_item),
-                        cells=[],
-                    )
-                    clusters.append(cluster)
-
-                # Map cells to clusters
-                # TODO: Remove, postprocess should take care of it anyway.
-                for cell in page.cells:
-                    for cluster in clusters:
-                        if not cell.bbox.area() > 0:
-                            overlap_frac = 0.0
-                        else:
-                            overlap_frac = (
-                                cell.bbox.intersection_area_with(cluster.bbox)
-                                / cell.bbox.area()
-                            )
-
-                        if overlap_frac > 0.5:
-                            cluster.cells.append(cell)
-
-                # Pre-sort clusters
-                # clusters = self.sort_clusters_by_cell_order(clusters)
-
-                # DEBUG code:
-                def draw_clusters_and_cells():
-                    image = copy.deepcopy(page.image)
-                    draw = ImageDraw.Draw(image)
-                    for c in clusters:
-                        x0, y0, x1, y1 = c.bbox.as_tuple()
-                        draw.rectangle([(x0, y0), (x1, y1)], outline="green")
-
-                        cell_color = (
-                            random.randint(30, 140),
-                            random.randint(30, 140),
-                            random.randint(30, 140),
+                with TimeRecorder(conv_res, "layout"):
+                    assert page.size is not None
+
+                    clusters = []
+                    for ix, pred_item in enumerate(
+                        self.layout_predictor.predict(page.get_image(scale=1.0))
+                    ):
+                        label = DocItemLabel(
+                            pred_item["label"]
+                            .lower()
+                            .replace(" ", "_")
+                            .replace("-", "_")
+                        )  # Temporary, until docling-ibm-model uses docling-core types
+                        cluster = Cluster(
+                            id=ix,
+                            label=label,
+                            confidence=pred_item["confidence"],
+                            bbox=BoundingBox.model_validate(pred_item),
+                            cells=[],
                         )
-                        for tc in c.cells:  # [:1]:
-                            x0, y0, x1, y1 = tc.bbox.as_tuple()
-                            draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
-                    image.show()
-
-                # draw_clusters_and_cells()
-
-                clusters, page.cells = self.postprocess(
-                    clusters, page.cells, page.size.height
-                )
+                        clusters.append(cluster)
+
+                    # Map cells to clusters
+                    # TODO: Remove, postprocess should take care of it anyway.
+                    for cell in page.cells:
+                        for cluster in clusters:
+                            if not cell.bbox.area() > 0:
+                                overlap_frac = 0.0
+                            else:
+                                overlap_frac = (
+                                    cell.bbox.intersection_area_with(cluster.bbox)
+                                    / cell.bbox.area()
+                                )
+
+                            if overlap_frac > 0.5:
+                                cluster.cells.append(cell)
+
+                    # Pre-sort clusters
+                    # clusters = self.sort_clusters_by_cell_order(clusters)
+
+                    # DEBUG code:
+                    def draw_clusters_and_cells(show: bool = False):
+                        image = copy.deepcopy(page.image)
+                        if image is not None:
+                            draw = ImageDraw.Draw(image)
+                            for c in clusters:
+                                x0, y0, x1, y1 = c.bbox.as_tuple()
+                                draw.rectangle([(x0, y0), (x1, y1)], outline="green")
+
+                                cell_color = (
+                                    random.randint(30, 140),
+                                    random.randint(30, 140),
+                                    random.randint(30, 140),
+                                )
+                                for tc in c.cells:  # [:1]:
+                                    x0, y0, x1, y1 = tc.bbox.as_tuple()
+                                    draw.rectangle(
+                                        [(x0, y0), (x1, y1)], outline=cell_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"layout_page_{page.page_no:05}.png"
+                                )
+                                image.save(str(out_file), format="png")
+
+                    # draw_clusters_and_cells()
+
+                    clusters, page.cells = self.postprocess(
+                        clusters, page.cells, page.size.height
+                    )
 
-                # draw_clusters_and_cells()
+                    page.predictions.layout = LayoutPrediction(clusters=clusters)
 
-                page.predictions.layout = LayoutPrediction(clusters=clusters)
+                if settings.debug.visualize_layout:
+                    draw_clusters_and_cells()
 
                 yield page

docling/models/page_assemble_model.py

@@ -12,8 +12,10 @@ from docling.datamodel.base_models import (
     Table,
     TextElement,
 )
+from docling.datamodel.document import ConversionResult
 from docling.models.base_model import BasePageModel
 from docling.models.layout_model import LayoutModel
+from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
 
@@ -51,122 +53,122 @@ class PageAssembleModel(BasePageModel):
 
         return sanitized_text.strip()  # Strip any leading or trailing whitespace
 
-    def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+    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:
-                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 == LayoutModel.TABLE_LABEL:
-                        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(
+                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="",
-                                otsl_seq=[],
-                                table_cells=[],
-                                cluster=cluster,
+                                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 == LayoutModel.TABLE_LABEL:
+                            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(
+                            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 == LayoutModel.FORMULA_LABEL:
-                        equation = None
-                        if page.predictions.equations_prediction:
-                            equation = (
-                                page.predictions.equations_prediction.equation_map.get(
+                            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 == LayoutModel.FORMULA_LABEL:
+                            equation = None
+                            if page.predictions.equations_prediction:
+                                equation = page.predictions.equations_prediction.equation_map.get(
                                     cluster.id, None
                                 )
-                            )
-                        if (
-                            not equation
-                        ):  # fallback: add empty formula, if it isn't present
-                            text = self.sanitize_text(
-                                [
-                                    cell.text.replace("\x02", "-").strip()
-                                    for cell in cluster.cells
-                                    if len(cell.text.strip()) > 0
-                                ]
-                            )
-                            equation = TextElement(
-                                label=cluster.label,
-                                id=cluster.id,
-                                cluster=cluster,
-                                page_no=page.page_no,
-                                text=text,
-                            )
-                        elements.append(equation)
-                        body.append(equation)
+                            if (
+                                not equation
+                            ):  # fallback: add empty formula, if it isn't present
+                                text = self.sanitize_text(
+                                    [
+                                        cell.text.replace("\x02", "-").strip()
+                                        for cell in cluster.cells
+                                        if len(cell.text.strip()) > 0
+                                    ]
+                                )
+                                equation = TextElement(
+                                    label=cluster.label,
+                                    id=cluster.id,
+                                    cluster=cluster,
+                                    page_no=page.page_no,
+                                    text=text,
+                                )
+                            elements.append(equation)
+                            body.append(equation)
 
-                page.assembled = AssembledUnit(
-                    elements=elements, headers=headers, body=body
-                )
+                    page.assembled = AssembledUnit(
+                        elements=elements, headers=headers, body=body
+                    )
 
-                # Remove page images (can be disabled)
-                if not self.options.keep_images:
-                    page._image_cache = {}
+                    # Remove page images (can be disabled)
+                    if not self.options.keep_images:
+                        page._image_cache = {}
 
-                # Unload backend
-                page._backend.unload()
+                    # Unload backend
+                    page._backend.unload()
 
                 yield page

docling/models/page_preprocessing_model.py

@@ -1,10 +1,14 @@
+from pathlib import Path
 from typing import Iterable, Optional
 
 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):
@@ -15,14 +19,17 @@ class PagePreprocessingModel(BasePageModel):
     def __init__(self, options: PagePreprocessingOptions):
         self.options = options
 
-    def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+    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:
-                page = self._populate_page_images(page)
-                page = self._parse_page_cells(page)
+                with TimeRecorder(conv_res, "page_parse"):
+                    page = self._populate_page_images(page)
+                    page = self._parse_page_cells(conv_res, page)
                 yield page
 
     # Generate the page image and store it in the page object
@@ -43,19 +50,30 @@ class PagePreprocessingModel(BasePageModel):
         return page
 
     # Extract and populate the page cells and store it in the page object
-    def _parse_page_cells(self, page: Page) -> Page:
+    def _parse_page_cells(self, conv_res: ConversionResult, page: Page) -> Page:
         assert page._backend is not None
 
         page.cells = list(page._backend.get_text_cells())
 
         # DEBUG code:
-        def draw_text_boxes(image, cells):
+        def draw_text_boxes(image, cells, show: bool = False):
             draw = ImageDraw.Draw(image)
             for c in cells:
                 x0, y0, x1, y1 = c.bbox.as_tuple()
                 draw.rectangle([(x0, y0), (x1, y1)], outline="red")
-            image.show()
+            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")
 
-        # draw_text_boxes(page.get_image(scale=1.0), cells)
+        if settings.debug.visualize_cells:
+            draw_text_boxes(page.get_image(scale=1.0), page.cells)
 
         return page