docling 2.11.0__py3-none-any.whl → 2.13.0__py3-none-any.whl

docling/models/layout_model.py

@@ -7,7 +7,7 @@ from typing import Iterable, List
 
 from docling_core.types.doc import CoordOrigin, DocItemLabel
 from docling_ibm_models.layoutmodel.layout_predictor import LayoutPredictor
-from PIL import ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 
 from docling.datamodel.base_models import (
     BoundingBox,
@@ -17,9 +17,11 @@ from docling.datamodel.base_models import (
     Page,
 )
 from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import AcceleratorDevice, AcceleratorOptions
 from docling.datamodel.settings import settings
 from docling.models.base_model import BasePageModel
-from docling.utils import layout_utils as lu
+from docling.utils.accelerator_utils import decide_device
+from docling.utils.layout_postprocessor import LayoutPostprocessor
 from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
@@ -42,237 +44,139 @@ class LayoutModel(BasePageModel):
     ]
     PAGE_HEADER_LABELS = [DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER]
 
-    TABLE_LABEL = DocItemLabel.TABLE
+    TABLE_LABELS = [DocItemLabel.TABLE, DocItemLabel.DOCUMENT_INDEX]
     FIGURE_LABEL = DocItemLabel.PICTURE
     FORMULA_LABEL = DocItemLabel.FORMULA
+    CONTAINER_LABELS = [DocItemLabel.FORM, DocItemLabel.KEY_VALUE_REGION]
 
-    def __init__(self, artifacts_path: Path):
-        self.layout_predictor = LayoutPredictor(artifacts_path)  # TODO temporary
+    def __init__(self, artifacts_path: Path, accelerator_options: AcceleratorOptions):
+        device = decide_device(accelerator_options.device)
 
-    def postprocess(self, clusters_in: List[Cluster], cells: List[Cell], page_height):
-        MIN_INTERSECTION = 0.2
-        CLASS_THRESHOLDS = {
-            DocItemLabel.CAPTION: 0.35,
-            DocItemLabel.FOOTNOTE: 0.35,
-            DocItemLabel.FORMULA: 0.35,
-            DocItemLabel.LIST_ITEM: 0.35,
-            DocItemLabel.PAGE_FOOTER: 0.35,
-            DocItemLabel.PAGE_HEADER: 0.35,
-            DocItemLabel.PICTURE: 0.2,  # low threshold adjust to capture chemical structures for examples.
-            DocItemLabel.SECTION_HEADER: 0.45,
-            DocItemLabel.TABLE: 0.35,
-            DocItemLabel.TEXT: 0.45,
-            DocItemLabel.TITLE: 0.45,
-            DocItemLabel.DOCUMENT_INDEX: 0.45,
-            DocItemLabel.CODE: 0.45,
-            DocItemLabel.CHECKBOX_SELECTED: 0.45,
-            DocItemLabel.CHECKBOX_UNSELECTED: 0.45,
-            DocItemLabel.FORM: 0.45,
-            DocItemLabel.KEY_VALUE_REGION: 0.45,
-        }
-
-        CLASS_REMAPPINGS = {
-            DocItemLabel.DOCUMENT_INDEX: DocItemLabel.TABLE,
-            DocItemLabel.TITLE: DocItemLabel.SECTION_HEADER,
-        }
-
-        _log.debug("================= Start postprocess function ====================")
-        start_time = time.time()
-        # Apply Confidence Threshold to cluster predictions
-        # confidence = self.conf_threshold
-        clusters_mod = []
-
-        for cluster in clusters_in:
-            confidence = CLASS_THRESHOLDS[cluster.label]
-            if cluster.confidence >= confidence:
-                # annotation["created_by"] = "high_conf_pred"
-
-                # Remap class labels where needed.
-                if cluster.label in CLASS_REMAPPINGS.keys():
-                    cluster.label = CLASS_REMAPPINGS[cluster.label]
-                clusters_mod.append(cluster)
-
-        # map to dictionary clusters and cells, with bottom left origin
-        clusters_orig = [
-            {
-                "id": c.id,
-                "bbox": list(
-                    c.bbox.to_bottom_left_origin(page_height).as_tuple()
-                ),  # TODO
-                "confidence": c.confidence,
-                "cell_ids": [],
-                "type": c.label,
-            }
-            for c in clusters_in
-        ]
-
-        clusters_out = [
-            {
-                "id": c.id,
-                "bbox": list(
-                    c.bbox.to_bottom_left_origin(page_height).as_tuple()
-                ),  # TODO
-                "confidence": c.confidence,
-                "created_by": "high_conf_pred",
-                "cell_ids": [],
-                "type": c.label,
-            }
-            for c in clusters_mod
-        ]
-
-        del clusters_mod
-
-        raw_cells = [
-            {
-                "id": c.id,
-                "bbox": list(
-                    c.bbox.to_bottom_left_origin(page_height).as_tuple()
-                ),  # TODO
-                "text": c.text,
-            }
-            for c in cells
-        ]
-        cell_count = len(raw_cells)
-
-        _log.debug("---- 0. Treat cluster overlaps ------")
-        clusters_out = lu.remove_cluster_duplicates_by_conf(clusters_out, 0.8)
-
-        _log.debug(
-            "---- 1. Initially assign cells to clusters based on minimum intersection ------"
-        )
-        ## Check for cells included in or touched by clusters:
-        clusters_out = lu.assigning_cell_ids_to_clusters(
-            clusters_out, raw_cells, MIN_INTERSECTION
+        self.layout_predictor = LayoutPredictor(
+            artifact_path=str(artifacts_path),
+            device=device,
+            num_threads=accelerator_options.num_threads,
         )
 
-        _log.debug("---- 2. Assign Orphans with Low Confidence Detections")
-        # Creates a map of cell_id->cluster_id
-        (
-            clusters_around_cells,
-            orphan_cell_indices,
-            ambiguous_cell_indices,
-        ) = lu.cell_id_state_map(clusters_out, cell_count)
-
-        # Assign orphan cells with lower confidence predictions
-        clusters_out, orphan_cell_indices = lu.assign_orphans_with_low_conf_pred(
-            clusters_out, clusters_orig, raw_cells, orphan_cell_indices
-        )
-
-        # Refresh the cell_ids assignment, after creating new clusters using low conf predictions
-        clusters_out = lu.assigning_cell_ids_to_clusters(
-            clusters_out, raw_cells, MIN_INTERSECTION
-        )
-
-        _log.debug("---- 3. Settle Ambigous Cells")
-        # Creates an update map after assignment of cell_id->cluster_id
-        (
-            clusters_around_cells,
-            orphan_cell_indices,
-            ambiguous_cell_indices,
-        ) = lu.cell_id_state_map(clusters_out, cell_count)
-
-        # Settle pdf cells that belong to multiple clusters
-        clusters_out, ambiguous_cell_indices = lu.remove_ambigous_pdf_cell_by_conf(
-            clusters_out, raw_cells, ambiguous_cell_indices
-        )
-
-        _log.debug("---- 4. Set Orphans as Text")
-        (
-            clusters_around_cells,
-            orphan_cell_indices,
-            ambiguous_cell_indices,
-        ) = lu.cell_id_state_map(clusters_out, cell_count)
-
-        clusters_out, orphan_cell_indices = lu.set_orphan_as_text(
-            clusters_out, clusters_orig, raw_cells, orphan_cell_indices
-        )
-
-        _log.debug("---- 5. Merge Cells & and adapt the bounding boxes")
-        # Merge cells orphan cells
-        clusters_out = lu.merge_cells(clusters_out)
-
-        # Clean up clusters that remain from merged and unreasonable clusters
-        clusters_out = lu.clean_up_clusters(
-            clusters_out,
-            raw_cells,
-            merge_cells=True,
-            img_table=True,
-            one_cell_table=True,
-        )
-
-        new_clusters = lu.adapt_bboxes(raw_cells, clusters_out, orphan_cell_indices)
-        clusters_out = new_clusters
-
-        ## We first rebuild where every cell is now:
-        ##   Now we write into a prediction cells list, not into the raw cells list.
-        ##   As we don't need previous labels, we best overwrite any old list, because that might
-        ##   have been sorted differently.
-        (
-            clusters_around_cells,
-            orphan_cell_indices,
-            ambiguous_cell_indices,
-        ) = lu.cell_id_state_map(clusters_out, cell_count)
-
-        target_cells = []
-        for ix, cell in enumerate(raw_cells):
-            new_cell = {
-                "id": ix,
-                "rawcell_id": ix,
-                "label": "None",
-                "bbox": cell["bbox"],
-                "text": cell["text"],
-            }
-            for cluster_index in clusters_around_cells[
-                ix
-            ]:  # By previous analysis, this is always 1 cluster.
-                new_cell["label"] = clusters_out[cluster_index]["type"]
-            target_cells.append(new_cell)
-            # _log.debug("New label of cell " + str(ix) + " is " + str(new_cell["label"]))
-        cells_out = target_cells
-
-        ## -------------------------------
-        ## Sort clusters into reasonable reading order, and sort the cells inside each cluster
-        _log.debug("---- 5. Sort clusters in reading order ------")
-        sorted_clusters = lu.produce_reading_order(
-            clusters_out, "raw_cell_ids", "raw_cell_ids", True
-        )
-        clusters_out = sorted_clusters
-
-        # end_time = timer()
-        _log.debug("---- End of postprocessing function ------")
-        end_time = time.time() - start_time
-        _log.debug(f"Finished post processing in seconds={end_time:.3f}")
-
-        cells_out_new = [
-            Cell(
-                id=c["id"],  # type: ignore
-                bbox=BoundingBox.from_tuple(
-                    coord=c["bbox"], origin=CoordOrigin.BOTTOMLEFT  # type: ignore
-                ).to_top_left_origin(page_height),
-                text=c["text"],  # type: ignore
-            )
-            for c in cells_out
-        ]
-
-        del cells_out
+    def draw_clusters_and_cells_side_by_side(
+        self, conv_res, page, clusters, mode_prefix: str, show: bool = False
+    ):
+        """
+        Draws a page image side by side with clusters filtered into two categories:
+        - Left: Clusters excluding FORM, KEY_VALUE_REGION, and PICTURE.
+        - Right: Clusters including FORM, KEY_VALUE_REGION, and PICTURE.
+        Includes label names and confidence scores for each cluster.
+        """
+        label_to_color = {
+            DocItemLabel.TEXT: (255, 255, 153),  # Light Yellow
+            DocItemLabel.CAPTION: (255, 204, 153),  # Light Orange
+            DocItemLabel.LIST_ITEM: (153, 153, 255),  # Light Purple
+            DocItemLabel.FORMULA: (192, 192, 192),  # Gray
+            DocItemLabel.TABLE: (255, 204, 204),  # Light Pink
+            DocItemLabel.PICTURE: (255, 204, 164),  # Light Beige
+            DocItemLabel.SECTION_HEADER: (255, 153, 153),  # Light Red
+            DocItemLabel.PAGE_HEADER: (204, 255, 204),  # Light Green
+            DocItemLabel.PAGE_FOOTER: (
+                204,
+                255,
+                204,
+            ),  # Light Green (same as Page-Header)
+            DocItemLabel.TITLE: (255, 153, 153),  # Light Red (same as Section-Header)
+            DocItemLabel.FOOTNOTE: (200, 200, 255),  # Light Blue
+            DocItemLabel.DOCUMENT_INDEX: (220, 220, 220),  # Light Gray
+            DocItemLabel.CODE: (125, 125, 125),  # Gray
+            DocItemLabel.CHECKBOX_SELECTED: (255, 182, 193),  # Pale Green
+            DocItemLabel.CHECKBOX_UNSELECTED: (255, 182, 193),  # Light Pink
+            DocItemLabel.FORM: (200, 255, 255),  # Light Cyan
+            DocItemLabel.KEY_VALUE_REGION: (183, 65, 14),  # Rusty orange
+        }
+        # Filter clusters for left and right images
+        exclude_labels = {
+            DocItemLabel.FORM,
+            DocItemLabel.KEY_VALUE_REGION,
+            DocItemLabel.PICTURE,
+        }
+        left_clusters = [c for c in clusters if c.label not in exclude_labels]
+        right_clusters = [c for c in clusters if c.label in exclude_labels]
+        # Create a deep copy of the original image for both sides
+        left_image = copy.deepcopy(page.image)
+        right_image = copy.deepcopy(page.image)
+
+        # Function to draw clusters on an image
+        def draw_clusters(image, clusters):
+            draw = ImageDraw.Draw(image, "RGBA")
+            # Create a smaller font for the labels
+            try:
+                font = ImageFont.truetype("arial.ttf", 12)
+            except OSError:
+                # Fallback to default font if arial is not available
+                font = ImageFont.load_default()
+            for c_tl in clusters:
+                all_clusters = [c_tl, *c_tl.children]
+                for c in all_clusters:
+                    # Draw cells first (underneath)
+                    cell_color = (0, 0, 0, 40)  # Transparent black for cells
+                    for tc in c.cells:
+                        cx0, cy0, cx1, cy1 = tc.bbox.as_tuple()
+                        draw.rectangle(
+                            [(cx0, cy0), (cx1, cy1)],
+                            outline=None,
+                            fill=cell_color,
+                        )
+                    # Draw cluster rectangle
+                    x0, y0, x1, y1 = c.bbox.as_tuple()
+                    cluster_fill_color = (*list(label_to_color.get(c.label)), 70)
+                    cluster_outline_color = (*list(label_to_color.get(c.label)), 255)
+                    draw.rectangle(
+                        [(x0, y0), (x1, y1)],
+                        outline=cluster_outline_color,
+                        fill=cluster_fill_color,
+                    )
+                    # Add label name and confidence
+                    label_text = f"{c.label.name} ({c.confidence:.2f})"
+                    # Create semi-transparent background for text
+                    text_bbox = draw.textbbox((x0, y0), label_text, font=font)
+                    text_bg_padding = 2
+                    draw.rectangle(
+                        [
+                            (
+                                text_bbox[0] - text_bg_padding,
+                                text_bbox[1] - text_bg_padding,
+                            ),
+                            (
+                                text_bbox[2] + text_bg_padding,
+                                text_bbox[3] + text_bg_padding,
+                            ),
+                        ],
+                        fill=(255, 255, 255, 180),  # Semi-transparent white
+                    )
+                    # Draw text
+                    draw.text(
+                        (x0, y0),
+                        label_text,
+                        fill=(0, 0, 0, 255),  # Solid black
+                        font=font,
+                    )
 
-        clusters_out_new = []
-        for c in clusters_out:
-            cluster_cells = [
-                ccell for ccell in cells_out_new if ccell.id in c["cell_ids"]  # type: ignore
-            ]
-            c_new = Cluster(
-                id=c["id"],  # type: ignore
-                bbox=BoundingBox.from_tuple(
-                    coord=c["bbox"], origin=CoordOrigin.BOTTOMLEFT  # type: ignore
-                ).to_top_left_origin(page_height),
-                confidence=c["confidence"],  # type: ignore
-                label=DocItemLabel(c["type"]),
-                cells=cluster_cells,
+        # Draw clusters on both images
+        draw_clusters(left_image, left_clusters)
+        draw_clusters(right_image, right_clusters)
+        # Combine the images side by side
+        combined_width = left_image.width * 2
+        combined_height = left_image.height
+        combined_image = Image.new("RGB", (combined_width, combined_height))
+        combined_image.paste(left_image, (0, 0))
+        combined_image.paste(right_image, (left_image.width, 0))
+        if show:
+            combined_image.show()
+        else:
+            out_path: Path = (
+                Path(settings.debug.debug_output_path)
+                / f"debug_{conv_res.input.file.stem}"
             )
-            clusters_out_new.append(c_new)
-
-        return clusters_out_new, cells_out_new
+            out_path.mkdir(parents=True, exist_ok=True)
+            out_file = out_path / f"{mode_prefix}_layout_page_{page.page_no:05}.png"
+            combined_image.save(str(out_file), format="png")
 
     def __call__(
         self, conv_res: ConversionResult, page_batch: Iterable[Page]
@@ -305,66 +209,26 @@ class LayoutModel(BasePageModel):
                         )
                         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)
+                    if settings.debug.visualize_raw_layout:
+                        self.draw_clusters_and_cells_side_by_side(
+                            conv_res, page, clusters, mode_prefix="raw"
+                        )
 
-                                out_file = (
-                                    out_path / f"layout_page_{page.page_no:05}.png"
-                                )
-                                image.save(str(out_file), format="png")
+                    # Apply postprocessing
 
-                    # draw_clusters_and_cells()
+                    processed_clusters, processed_cells = LayoutPostprocessor(
+                        page.cells, clusters, page.size
+                    ).postprocess()
+                    # processed_clusters, processed_cells = clusters, page.cells
 
-                    clusters, page.cells = self.postprocess(
-                        clusters, page.cells, page.size.height
+                    page.cells = processed_cells
+                    page.predictions.layout = LayoutPrediction(
+                        clusters=processed_clusters
                     )
 
-                    page.predictions.layout = LayoutPrediction(clusters=clusters)
-
                 if settings.debug.visualize_layout:
-                    draw_clusters_and_cells()
+                    self.draw_clusters_and_cells_side_by_side(
+                        conv_res, page, processed_clusters, mode_prefix="postprocessed"
+                    )
 
                 yield page

docling/models/page_assemble_model.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel
 
 from docling.datamodel.base_models import (
     AssembledUnit,
+    ContainerElement,
     FigureElement,
     Page,
     PageElement,
@@ -94,7 +95,7 @@ class PageAssembleModel(BasePageModel):
                                 headers.append(text_el)
                             else:
                                 body.append(text_el)
-                        elif cluster.label == LayoutModel.TABLE_LABEL:
+                        elif cluster.label in LayoutModel.TABLE_LABELS:
                             tbl = None
                             if page.predictions.tablestructure:
                                 tbl = page.predictions.tablestructure.table_map.get(
@@ -159,6 +160,15 @@ class PageAssembleModel(BasePageModel):
                                 )
                             elements.append(equation)
                             body.append(equation)
+                        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

docling/models/rapid_ocr_model.py

@@ -6,16 +6,26 @@ 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 RapidOcrOptions
+from docling.datamodel.pipeline_options import (
+    AcceleratorDevice,
+    AcceleratorOptions,
+    RapidOcrOptions,
+)
 from docling.datamodel.settings import settings
 from docling.models.base_ocr_model import BaseOcrModel
+from docling.utils.accelerator_utils import decide_device
 from docling.utils.profiling import TimeRecorder
 
 _log = logging.getLogger(__name__)
 
 
 class RapidOcrModel(BaseOcrModel):
-    def __init__(self, enabled: bool, options: RapidOcrOptions):
+    def __init__(
+        self,
+        enabled: bool,
+        options: RapidOcrOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
         super().__init__(enabled=enabled, options=options)
         self.options: RapidOcrOptions
 
@@ -30,52 +40,21 @@ class RapidOcrModel(BaseOcrModel):
                     "Alternatively, Docling has support for other OCR engines. See the documentation."
                 )
 
-            # This configuration option will be revamped while introducing device settings for all models.
-            # For the moment we will default to auto and let onnx-runtime pick the best.
-            cls_use_cuda = True
-            rec_use_cuda = True
-            det_use_cuda = True
-            det_use_dml = True
-            cls_use_dml = True
-            rec_use_dml = True
-
-            # # Same as Defaults in RapidOCR
-            # cls_use_cuda = False
-            # rec_use_cuda = False
-            # det_use_cuda = False
-            # det_use_dml = False
-            # cls_use_dml = False
-            # rec_use_dml = False
-
-            # # If we set everything to true onnx-runtime would automatically choose the fastest accelerator
-            # if self.options.device == self.options.Device.AUTO:
-            #     cls_use_cuda = True
-            #     rec_use_cuda = True
-            #     det_use_cuda = True
-            #     det_use_dml = True
-            #     cls_use_dml = True
-            #     rec_use_dml = True
-
-            # # If we set use_cuda to true onnx would use the cuda device available in runtime if no cuda device is available it would run on CPU.
-            # elif self.options.device == self.options.Device.CUDA:
-            #     cls_use_cuda = True
-            #     rec_use_cuda = True
-            #     det_use_cuda = True
-
-            # # If we set use_dml to true onnx would use the dml device available in runtime if no dml device is available it would work on CPU.
-            # elif self.options.device == self.options.Device.DIRECTML:
-            #     det_use_dml = True
-            #     cls_use_dml = True
-            #     rec_use_dml = True
+            # Decide the accelerator devices
+            device = decide_device(accelerator_options.device)
+            use_cuda = str(AcceleratorDevice.CUDA.value).lower() in device
+            use_dml = accelerator_options.device == AcceleratorDevice.AUTO
+            intra_op_num_threads = accelerator_options.num_threads
 
             self.reader = RapidOCR(
                 text_score=self.options.text_score,
-                cls_use_cuda=cls_use_cuda,
-                rec_use_cuda=rec_use_cuda,
-                det_use_cuda=det_use_cuda,
-                det_use_dml=det_use_dml,
-                cls_use_dml=cls_use_dml,
-                rec_use_dml=rec_use_dml,
+                cls_use_cuda=use_cuda,
+                rec_use_cuda=use_cuda,
+                det_use_cuda=use_cuda,
+                det_use_dml=use_dml,
+                cls_use_dml=use_dml,
+                rec_use_dml=use_dml,
+                intra_op_num_threads=intra_op_num_threads,
                 print_verbose=self.options.print_verbose,
                 det_model_path=self.options.det_model_path,
                 cls_model_path=self.options.cls_model_path,