docling-core 2.43.1__py3-none-any.whl → 2.44.1__py3-none-any.whl

docling_core/transforms/visualizer/key_value_visualizer.py

@@ -0,0 +1,217 @@
+"""Key‑value visualizer overlaying key/value cells and their links on page images.
+
+This module complements :py:class:`layout_visualizer.LayoutVisualizer` by drawing
+*key* and *value* cells plus the directed links between them.  It can be stacked
+on top of any other :py:class:`BaseVisualizer` – e.g. first draw the general
+layout, then add the key‑value layer.
+"""
+
+from copy import deepcopy
+from typing import Optional, Union
+
+from PIL import ImageDraw, ImageFont
+from PIL.Image import Image
+from PIL.ImageFont import FreeTypeFont
+from pydantic import BaseModel
+from typing_extensions import override
+
+from docling_core.transforms.visualizer.base import BaseVisualizer
+from docling_core.types.doc.document import ContentLayer, DoclingDocument
+from docling_core.types.doc.labels import GraphCellLabel, GraphLinkLabel
+
+# ---------------------------------------------------------------------------
+# Helper functions / constants
+# ---------------------------------------------------------------------------
+
+# Semi‑transparent RGBA colours for key / value cells and their connecting link
+_KEY_FILL = (0, 170, 0, 70)  # greenish
+_VALUE_FILL = (0, 0, 200, 70)  # bluish
+_LINK_COLOUR = (255, 0, 0, 255)  # red line (solid)
+
+_LABEL_TXT_COLOUR = (0, 0, 0, 255)
+_LABEL_BG_COLOUR = (255, 255, 255, 180)  # semi‑transparent white
+
+
+class KeyValueVisualizer(BaseVisualizer):
+    """Draw key/value graphs stored in :py:attr:`DoclingDocument.key_value_items`."""
+
+    class Params(BaseModel):
+        """Parameters for KeyValueVisualizer controlling label and cell id display, and content layers to visualize."""
+
+        show_label: bool = True  # draw cell text close to bbox
+        show_cell_id: bool = False  # annotate each rectangle with its cell_id
+        content_layers: set[ContentLayer] = {cl for cl in ContentLayer}
+
+    base_visualizer: Optional[BaseVisualizer] = None
+    params: Params = Params()
+
+    # ---------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------
+
+    def _cell_fill(self, label: GraphCellLabel) -> tuple[int, int, int, int]:
+        """Return RGBA fill colour depending on *label*."""
+        return _KEY_FILL if label == GraphCellLabel.KEY else _VALUE_FILL
+
+    def _draw_key_value_layer(
+        self,
+        *,
+        image: Image,
+        doc: DoclingDocument,
+        page_no: int,
+        scale_x: float,
+        scale_y: float,
+    ) -> None:
+        """Draw every key‑value graph that has cells on *page_no* onto *image*."""
+        draw = ImageDraw.Draw(image, "RGBA")
+        # Choose a small truetype font if available, otherwise default bitmap font
+        font: Union[ImageFont.ImageFont, FreeTypeFont]
+        try:
+            font = ImageFont.truetype("arial.ttf", 12)
+        except OSError:
+            font = ImageFont.load_default()
+
+        for kv_item in doc.key_value_items:
+            cell_dict = {cell.cell_id: cell for cell in kv_item.graph.cells}
+
+            # ------------------------------------------------------------------
+            # First draw cells (rectangles + optional labels)
+            # ------------------------------------------------------------------
+            for cell in cell_dict.values():
+                if cell.prov is None or cell.prov.page_no != page_no:
+                    continue  # skip cells not on this page or without bbox
+
+                tl_bbox = cell.prov.bbox.to_top_left_origin(
+                    page_height=doc.pages[page_no].size.height
+                )
+                x0, y0, x1, y1 = tl_bbox.as_tuple()
+                x0 *= scale_x
+                x1 *= scale_x
+                y0 *= scale_y
+                y1 *= scale_y
+                fill_rgba = self._cell_fill(cell.label)
+
+                draw.rectangle(
+                    [(x0, y0), (x1, y1)],
+                    outline=fill_rgba[:-1] + (255,),
+                    fill=fill_rgba,
+                )
+
+                if self.params.show_label:
+                    txt_parts = []
+                    if self.params.show_cell_id:
+                        txt_parts.append(str(cell.cell_id))
+                    txt_parts.append(cell.text)
+                    label_text = " | ".join(txt_parts)
+
+                    tbx = draw.textbbox((x0, y0), label_text, font=font)
+                    pad = 2
+                    draw.rectangle(
+                        [(tbx[0] - pad, tbx[1] - pad), (tbx[2] + pad, tbx[3] + pad)],
+                        fill=_LABEL_BG_COLOUR,
+                    )
+                    draw.text((x0, y0), label_text, font=font, fill=_LABEL_TXT_COLOUR)
+
+            # ------------------------------------------------------------------
+            # Then draw links (after rectangles so they appear on top)
+            # ------------------------------------------------------------------
+            for link in kv_item.graph.links:
+                if link.label != GraphLinkLabel.TO_VALUE:
+                    # Future‑proof: ignore other link types silently
+                    continue
+
+                src_cell = cell_dict.get(link.source_cell_id)
+                tgt_cell = cell_dict.get(link.target_cell_id)
+                if src_cell is None or tgt_cell is None:
+                    continue
+                if (
+                    src_cell.prov is None
+                    or tgt_cell.prov is None
+                    or src_cell.prov.page_no != page_no
+                    or tgt_cell.prov.page_no != page_no
+                ):
+                    continue  # only draw if both ends are on this page
+
+                def _centre(bbox):
+                    tl = bbox.to_top_left_origin(
+                        page_height=doc.pages[page_no].size.height
+                    )
+                    l, t, r, b = tl.as_tuple()
+                    return ((l + r) / 2 * scale_x, (t + b) / 2 * scale_y)
+
+                src_xy = _centre(src_cell.prov.bbox)
+                tgt_xy = _centre(tgt_cell.prov.bbox)
+
+                draw.line([src_xy, tgt_xy], fill=_LINK_COLOUR, width=2)
+
+                # draw a small arrow‑head by rendering a short orthogonal line
+                # segment; exact geometry is not critical for visual inspection
+                arrow_len = 6
+                dx = tgt_xy[0] - src_xy[0]
+                dy = tgt_xy[1] - src_xy[1]
+                length = (dx**2 + dy**2) ** 0.5 or 1.0
+                ux, uy = dx / length, dy / length
+                # perpendicular vector
+                px, py = -uy, ux
+                # two points forming the arrow head triangle base
+                head_base_left = (
+                    tgt_xy[0] - ux * arrow_len - px * arrow_len / 2,
+                    tgt_xy[1] - uy * arrow_len - py * arrow_len / 2,
+                )
+                head_base_right = (
+                    tgt_xy[0] - ux * arrow_len + px * arrow_len / 2,
+                    tgt_xy[1] - uy * arrow_len + py * arrow_len / 2,
+                )
+                draw.polygon(
+                    [tgt_xy, head_base_left, head_base_right], fill=_LINK_COLOUR
+                )
+
+    # ---------------------------------------------------------------------
+    # Public API – BaseVisualizer implementation
+    # ---------------------------------------------------------------------
+
+    @override
+    def get_visualization(
+        self,
+        *,
+        doc: DoclingDocument,
+        included_content_layers: Optional[set[ContentLayer]] = None,
+        **kwargs,
+    ) -> dict[Optional[int], Image]:
+        """Return page‑wise images with key/value overlay (incl. base layer)."""
+        base_images = (
+            self.base_visualizer.get_visualization(
+                doc=doc, included_content_layers=included_content_layers, **kwargs
+            )
+            if self.base_visualizer
+            else None
+        )
+
+        if included_content_layers is None:
+            included_content_layers = {cl for cl in ContentLayer}
+
+        images: dict[Optional[int], Image] = {}
+
+        # Ensure we have page images to draw on
+        for page_nr, page in doc.pages.items():
+            base_img = (base_images or {}).get(page_nr)
+            if base_img is None:
+                if page.image is None or (pil_img := page.image.pil_image) is None:
+                    raise RuntimeError("Cannot visualize document without page images")
+                base_img = deepcopy(pil_img)
+            images[page_nr] = base_img
+
+        # Overlay key‑value content
+        for page_nr, img in images.items():  # type: ignore
+            assert isinstance(page_nr, int)
+            scale_x = img.width / doc.pages[page_nr].size.width
+            scale_y = img.height / doc.pages[page_nr].size.height
+            self._draw_key_value_layer(
+                image=img,
+                doc=doc,
+                page_no=page_nr,
+                scale_x=scale_x,
+                scale_y=scale_y,
+            )
+
+        return images

docling_core/types/doc/document.py

@@ -1373,11 +1373,12 @@ class PictureItem(FloatingItem):
         )  # Encode to Base64 and decode to string
         return img_base64
 
-    def _image_to_hexhash(self) -> Optional[str]:
+    @staticmethod
+    def _image_to_hexhash(img: Optional[PILImage.Image]) -> Optional[str]:
         """Hexash from the image."""
-        if self.image is not None and self.image._pil is not None:
+        if img is not None:
             # Convert the image to raw bytes
-            image_bytes = self.image._pil.tobytes()
+            image_bytes = img.tobytes()
 
             # Create a hash object (e.g., SHA-256)
             hasher = hashlib.sha256(usedforsecurity=False)
@@ -4116,16 +4117,10 @@ class DoclingDocument(BaseModel):
         if image_dir.is_dir():
             for item, level in result.iterate_items(page_no=page_no, with_groups=False):
                 if isinstance(item, PictureItem):
+                    img = item.get_image(doc=self)
+                    if img is not None:
 
-                    if (
-                        item.image is not None
-                        and isinstance(item.image.uri, AnyUrl)
-                        and item.image.uri.scheme == "data"
-                        and item.image.pil_image is not None
-                    ):
-                        img = item.image.pil_image
-
-                        hexhash = item._image_to_hexhash()
+                        hexhash = PictureItem._image_to_hexhash(img)
 
                         # loc_path = image_dir / f"image_{img_count:06}.png"
                         if hexhash is not None:
@@ -4140,6 +4135,11 @@ class DoclingDocument(BaseModel):
                             else:
                                 obj_path = loc_path
 
+                            if item.image is None:
+                                scale = img.size[0] / item.prov[0].bbox.width
+                                item.image = ImageRef.from_pil(
+                                    image=img, dpi=round(72 * scale)
+                                )
                             item.image.uri = Path(obj_path)
 
                         # if item.image._pil is not None:
@@ -4539,6 +4539,8 @@ class DoclingDocument(BaseModel):
             reference_path = None
         else:
             reference_path = filename.parent
+            artifacts_dir = reference_path / artifacts_dir
+
         return artifacts_dir, reference_path
 
     def _make_copy_with_refmode(
@@ -5543,8 +5545,27 @@ class DoclingDocument(BaseModel):
         self,
         show_label: bool = True,
         show_branch_numbering: bool = False,
+        viz_mode: Literal["reading_order", "key_value"] = "reading_order",
+        show_cell_id: bool = False,
     ) -> dict[Optional[int], PILImage.Image]:
-        """Get visualization of the document as images by page."""
+        """Get visualization of the document as images by page.
+
+        :param show_label: Show labels on elements (applies to all visualizers).
+        :type show_label: bool
+        :param show_branch_numbering: Show branch numbering (reading order visualizer only).
+        :type show_branch_numbering: bool
+        :param visualizer: Which visualizer to use. One of 'reading_order' (default), 'key_value'.
+        :type visualizer: str
+        :param show_cell_id: Show cell IDs (key value visualizer only).
+        :type show_cell_id: bool
+
+        :returns: Dictionary mapping page numbers to PIL images.
+        :rtype: dict[Optional[int], PILImage.Image]
+        """
+        from docling_core.transforms.visualizer.base import BaseVisualizer
+        from docling_core.transforms.visualizer.key_value_visualizer import (
+            KeyValueVisualizer,
+        )
         from docling_core.transforms.visualizer.layout_visualizer import (
             LayoutVisualizer,
         )
@@ -5552,18 +5573,34 @@ class DoclingDocument(BaseModel):
             ReadingOrderVisualizer,
         )
 
-        visualizer = ReadingOrderVisualizer(
-            base_visualizer=LayoutVisualizer(
-                params=LayoutVisualizer.Params(
+        visualizer_obj: BaseVisualizer
+        if viz_mode == "reading_order":
+            visualizer_obj = ReadingOrderVisualizer(
+                base_visualizer=LayoutVisualizer(
+                    params=LayoutVisualizer.Params(
+                        show_label=show_label,
+                    ),
+                ),
+                params=ReadingOrderVisualizer.Params(
+                    show_branch_numbering=show_branch_numbering,
+                ),
+            )
+        elif viz_mode == "key_value":
+            visualizer_obj = KeyValueVisualizer(
+                base_visualizer=LayoutVisualizer(
+                    params=LayoutVisualizer.Params(
+                        show_label=show_label,
+                    ),
+                ),
+                params=KeyValueVisualizer.Params(
                     show_label=show_label,
+                    show_cell_id=show_cell_id,
                 ),
-            ),
-            params=ReadingOrderVisualizer.Params(
-                show_branch_numbering=show_branch_numbering,
-            ),
-        )
-        images = visualizer.get_visualization(doc=self)
+            )
+        else:
+            raise ValueError(f"Unknown visualization mode: {viz_mode}")
 
+        images = visualizer_obj.get_visualization(doc=self)
         return images
 
     @field_validator("version")

{docling_core-2.43.1.dist-info → docling_core-2.44.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docling-core
-Version: 2.43.1
+Version: 2.44.1
 Summary: A python library to define and validate data types in Docling.
 Author-email: Cesar Berrospi Ramis <ceb@zurich.ibm.com>, Panos Vagenas <pva@zurich.ibm.com>, Michele Dolfi <dol@zurich.ibm.com>, Christoph Auer <cau@zurich.ibm.com>, Peter Staar <taa@zurich.ibm.com>
 Maintainer-email: Panos Vagenas <pva@zurich.ibm.com>, Michele Dolfi <dol@zurich.ibm.com>, Christoph Auer <cau@zurich.ibm.com>, Peter Staar <taa@zurich.ibm.com>, Cesar Berrospi Ramis <ceb@zurich.ibm.com>

{docling_core-2.43.1.dist-info → docling_core-2.44.1.dist-info}/RECORD

@@ -35,6 +35,7 @@ docling_core/transforms/serializer/html_styles.py,sha256=-jBwS4EU7yfKoz0GSoxhwx9
 docling_core/transforms/serializer/markdown.py,sha256=VwonuAkuOPmQM7ibDIGvQBHOqhTcTJ_t187fLQQiNPo,23951
 docling_core/transforms/visualizer/__init__.py,sha256=gUfF25yiJ_KO46ZIUNqZQOZGy2PLx6gnnr6AZYxKHXI,35
 docling_core/transforms/visualizer/base.py,sha256=aEF7b3rHq6DVdX8zDYEPoq55BHDYe4Hh_97lBdcW4lY,555
+docling_core/transforms/visualizer/key_value_visualizer.py,sha256=fp7nFLy4flOSiavdRgg5y1Mu7WVLIDGh1zEHsq8kgVM,8979
 docling_core/transforms/visualizer/layout_visualizer.py,sha256=k93ORWxA1oTDNkNxyqlzqB2aunxBX3aq4e4ihdJzRkE,8089
 docling_core/transforms/visualizer/reading_order_visualizer.py,sha256=muqmaxOBao39X3Dut0934NAjU3I4v3JN5VzzdjmoGRY,7776
 docling_core/transforms/visualizer/table_visualizer.py,sha256=iJPjk-XQSSCH3oujcjPMz-redAwNNHseZ41lFyd-u3k,8097
@@ -42,7 +43,7 @@ docling_core/types/__init__.py,sha256=MVRSgsk5focwGyAplh_TRR3dEecIXpd98g_u3zZ5HX
 docling_core/types/base.py,sha256=PusJskRVL19y-hq0BgXr5e8--QEqSqLnFNJ8UbOqW88,8318
 docling_core/types/doc/__init__.py,sha256=8hOhm5W9mArf3zwgfoMxDs1pHizhLFSAZlLu1tPBBRk,1641
 docling_core/types/doc/base.py,sha256=i98y4IF250adR-8BSS374K90fwfwG-vBfWh14tLC5Cs,15906
-docling_core/types/doc/document.py,sha256=SUqIJ-huO3ELLRdCMUYjkkXHeGGkeY2oaOWFQ1nq5lg,199315
+docling_core/types/doc/document.py,sha256=-cL4eGFRbQHgXAsCG8zALxAx-IoanvkqG5E1zvKOMxI,201012
 docling_core/types/doc/labels.py,sha256=-W1-LW6z0J9F9ExJqR0Wd1WeqWTaY3Unm-j1UkQGlC4,7330
 docling_core/types/doc/page.py,sha256=35h1xdtCM3-AaN8Dim9jDseZIiw-3GxpB-ofF-H2rQQ,41878
 docling_core/types/doc/tokens.py,sha256=z22l9J81_sg9CYMvOuLmPuLsNT7h_s7wao2UT89DvI8,9278
@@ -75,9 +76,9 @@ docling_core/utils/generate_jsonschema.py,sha256=uNX1O5XnjyB5nA66XqZXTt3YbGuR2ty
 docling_core/utils/legacy.py,sha256=5lghO48OEcV9V51tRnH3YSKgLtdqhr-Q5C_OcJZ8TOs,24392
 docling_core/utils/validate.py,sha256=aQ11UbFyl8iD_N7yTTZmm_VVeXz8KcCyn3GLXgkfYRM,2049
 docling_core/utils/validators.py,sha256=azcrndLzhNkTWnbFSu9shJ5D3j_znnLrIFA5R8hzmGU,2798
-docling_core-2.43.1.dist-info/licenses/LICENSE,sha256=2M9-6EoQ1sxFztTOkXGAtwUDJvnWaAHdB9BYWVwGkIw,1087
-docling_core-2.43.1.dist-info/METADATA,sha256=uVJIJpT7DuKgWIYmdqxgTf2Av_5dB2Da9ZfcPiHqbPc,6453
-docling_core-2.43.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-docling_core-2.43.1.dist-info/entry_points.txt,sha256=ER4zROQWkFMHIrY-oqY5E4HeCcCIg8dLkNztYGxdb7c,59
-docling_core-2.43.1.dist-info/top_level.txt,sha256=O-tcXpGiurlud-1ZxMq1b-OmrfAVA4sajcgWU32RtfA,13
-docling_core-2.43.1.dist-info/RECORD,,
+docling_core-2.44.1.dist-info/licenses/LICENSE,sha256=2M9-6EoQ1sxFztTOkXGAtwUDJvnWaAHdB9BYWVwGkIw,1087
+docling_core-2.44.1.dist-info/METADATA,sha256=NtybqGuK3bjSWq_AadW7B-pydco80WLkFYwWaWyaTb0,6453
+docling_core-2.44.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+docling_core-2.44.1.dist-info/entry_points.txt,sha256=ER4zROQWkFMHIrY-oqY5E4HeCcCIg8dLkNztYGxdb7c,59
+docling_core-2.44.1.dist-info/top_level.txt,sha256=O-tcXpGiurlud-1ZxMq1b-OmrfAVA4sajcgWU32RtfA,13
+docling_core-2.44.1.dist-info/RECORD,,