docling 2.69.0__py3-none-any.whl

docling/utils/export.py

@@ -0,0 +1,146 @@
+import logging
+from collections.abc import Iterable
+from typing import Any, Dict, List, Tuple, Union
+
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.legacy_doc.base import BaseCell, BaseText, Ref, Table
+
+from docling.datamodel.document import ConversionResult, Page
+
+_log = logging.getLogger(__name__)
+
+
+def generate_multimodal_pages(
+    doc_result: ConversionResult,
+) -> Iterable[Tuple[str, str, List[Dict[str, Any]], List[Dict[str, Any]], Page]]:
+    label_to_doclaynet = {
+        "title": "title",
+        "table-of-contents": "document_index",
+        "subtitle-level-1": "section_header",
+        "checkbox-selected": "checkbox_selected",
+        "checkbox-unselected": "checkbox_unselected",
+        "caption": "caption",
+        "page-header": "page_header",
+        "page-footer": "page_footer",
+        "footnote": "footnote",
+        "table": "table",
+        "formula": "formula",
+        "list-item": "list_item",
+        "code": "code",
+        "figure": "picture",
+        "picture": "picture",
+        "reference": "text",
+        "paragraph": "text",
+        "text": "text",
+    }
+
+    content_text = ""
+    page_no = 0
+    start_ix = 0
+    end_ix = 0
+    doc_items: List[Tuple[int, Union[BaseCell, BaseText]]] = []
+
+    doc = doc_result.legacy_document
+
+    def _process_page_segments(doc_items: list[Tuple[int, BaseCell]], page: Page):
+        segments = []
+
+        for ix, item in doc_items:
+            item_type = item.obj_type
+            label = label_to_doclaynet.get(item_type, None)
+
+            if label is None or item.prov is None or page.size is None:
+                continue
+
+            bbox = BoundingBox.from_tuple(
+                tuple(item.prov[0].bbox), origin=CoordOrigin.BOTTOMLEFT
+            )
+            new_bbox = bbox.to_top_left_origin(page_height=page.size.height).normalized(
+                page_size=page.size
+            )
+
+            new_segment = {
+                "index_in_doc": ix,
+                "label": label,
+                "text": item.text if item.text is not None else "",
+                "bbox": new_bbox.as_tuple(),
+                "data": [],
+            }
+
+            if isinstance(item, Table):
+                table_html = item.export_to_html()
+                new_segment["data"].append(
+                    {
+                        "html_seq": table_html,
+                        "otsl_seq": "",
+                    }
+                )
+
+            segments.append(new_segment)
+
+        return segments
+
+    def _process_page_cells(page: Page):
+        cells: List[dict] = []
+        if page.size is None:
+            return cells
+        for cell in page.cells:
+            new_bbox = (
+                cell.rect.to_bounding_box()
+                .to_top_left_origin(page_height=page.size.height)
+                .normalized(page_size=page.size)
+            )
+            is_ocr = cell.from_ocr
+            ocr_confidence = cell.confidence
+            cells.append(
+                {
+                    "text": cell.text,
+                    "bbox": new_bbox.as_tuple(),
+                    "ocr": is_ocr,
+                    "ocr_confidence": ocr_confidence,
+                }
+            )
+        return cells
+
+    def _process_page():
+        page_ix = page_no - 1
+        page = doc_result.pages[page_ix]
+
+        page_cells = _process_page_cells(page=page)
+        page_segments = _process_page_segments(doc_items=doc_items, page=page)
+        content_md = doc.export_to_markdown(
+            main_text_start=start_ix, main_text_stop=end_ix
+        )
+        # No page-tagging since we only do 1 page at the time
+        content_dt = doc.export_to_document_tokens(
+            main_text_start=start_ix, main_text_stop=end_ix, add_page_index=False
+        )
+
+        return content_text, content_md, content_dt, page_cells, page_segments, page
+
+    if doc.main_text is None:
+        return
+    for ix, orig_item in enumerate(doc.main_text):
+        item = doc._resolve_ref(orig_item) if isinstance(orig_item, Ref) else orig_item
+        if item is None or item.prov is None or len(item.prov) == 0:
+            _log.debug(f"Skipping item {orig_item}")
+            continue
+
+        item_page = item.prov[0].page
+
+        # Page is complete
+        if page_no > 0 and item_page > page_no:
+            yield _process_page()
+
+            start_ix = ix
+            doc_items = []
+            content_text = ""
+
+        page_no = item_page
+        end_ix = ix
+        doc_items.append((ix, item))
+        if item.text is not None and item.text != "":
+            content_text += item.text + " "
+
+    if len(doc_items) > 0:
+        yield _process_page()

docling/utils/glm_utils.py

@@ -0,0 +1,361 @@
+import re
+from pathlib import Path
+from typing import List
+
+import pandas as pd
+from docling_core.types.doc import (
+    BoundingBox,
+    CoordOrigin,
+    DocItemLabel,
+    DoclingDocument,
+    DocumentOrigin,
+    GroupLabel,
+    ProvenanceItem,
+    Size,
+    TableCell,
+    TableData,
+)
+from docling_core.types.doc.document import ContentLayer
+
+
+def resolve_item(paths, obj):
+    """Find item in document from a reference path"""
+
+    if len(paths) == 0:
+        return obj
+
+    if paths[0] == "#":
+        return resolve_item(paths[1:], obj)
+
+    try:
+        key = int(paths[0])
+    except Exception:
+        key = paths[0]
+
+    if len(paths) == 1:
+        if isinstance(key, str) and key in obj:
+            return obj[key]
+        elif isinstance(key, int) and key < len(obj):
+            return obj[key]
+        else:
+            return None
+
+    elif len(paths) > 1:
+        if isinstance(key, str) and key in obj:
+            return resolve_item(paths[1:], obj[key])
+        elif isinstance(key, int) and key < len(obj):
+            return resolve_item(paths[1:], obj[key])
+        else:
+            return None
+
+    else:
+        return None
+
+
+def _flatten_table_grid(grid: List[List[dict]]) -> List[dict]:
+    unique_objects = []
+    seen_spans = set()
+
+    for sublist in grid:
+        for obj in sublist:
+            # Convert the spans list to a tuple of tuples for hashing
+            spans_tuple = tuple(tuple(span) for span in obj["spans"])
+            if spans_tuple not in seen_spans:
+                seen_spans.add(spans_tuple)
+                unique_objects.append(obj)
+
+    return unique_objects
+
+
+def to_docling_document(doc_glm, update_name_label=False) -> DoclingDocument:  # noqa: C901
+    origin = DocumentOrigin(
+        mimetype="application/pdf",
+        filename=doc_glm["file-info"]["filename"],
+        binary_hash=doc_glm["file-info"]["document-hash"],
+    )
+    doc_name = Path(origin.filename).stem
+
+    doc: DoclingDocument = DoclingDocument(name=doc_name, origin=origin)
+
+    for page_dim in doc_glm["page-dimensions"]:
+        page_no = int(page_dim["page"])
+        size = Size(width=page_dim["width"], height=page_dim["height"])
+
+        doc.add_page(page_no=page_no, size=size)
+
+    if "properties" in doc_glm:
+        props = pd.DataFrame(
+            doc_glm["properties"]["data"], columns=doc_glm["properties"]["headers"]
+        )
+    else:
+        props = pd.DataFrame()
+
+    current_list = None
+
+    for ix, pelem in enumerate(doc_glm["page-elements"]):
+        ptype = pelem["type"]
+        span_i = pelem["span"][0]
+        span_j = pelem["span"][1]
+
+        if "iref" not in pelem:
+            # print(json.dumps(pelem, indent=2))
+            continue
+
+        iref = pelem["iref"]
+
+        if re.match("#/figures/(\\d+)/captions/(.+)", iref):
+            # print(f"skip {iref}")
+            continue
+
+        if re.match("#/tables/(\\d+)/captions/(.+)", iref):
+            # print(f"skip {iref}")
+            continue
+
+        path = iref.split("/")
+        obj = resolve_item(path, doc_glm)
+
+        if obj is None:
+            current_list = None
+            print(f"warning: undefined {path}")
+            continue
+
+        if ptype == "figure":
+            current_list = None
+            text = ""
+            caption_refs = []
+            for caption in obj["captions"]:
+                text += caption["text"]
+
+                for nprov in caption["prov"]:
+                    npaths = nprov["$ref"].split("/")
+                    nelem = resolve_item(npaths, doc_glm)
+
+                    if nelem is None:
+                        # print(f"warning: undefined caption {npaths}")
+                        continue
+
+                    span_i = nelem["span"][0]
+                    span_j = nelem["span"][1]
+
+                    cap_text = caption["text"][span_i:span_j]
+
+                    # doc_glm["page-elements"].remove(nelem)
+
+                    prov = ProvenanceItem(
+                        page_no=nelem["page"],
+                        charspan=tuple(nelem["span"]),
+                        bbox=BoundingBox.from_tuple(
+                            nelem["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                        ),
+                    )
+
+                    caption_obj = doc.add_text(
+                        label=DocItemLabel.CAPTION, text=cap_text, prov=prov
+                    )
+                    caption_refs.append(caption_obj.get_ref())
+
+            prov = ProvenanceItem(
+                page_no=pelem["page"],
+                charspan=(0, len(text)),
+                bbox=BoundingBox.from_tuple(
+                    pelem["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                ),
+            )
+
+            pic = doc.add_picture(prov=prov)
+            pic.captions.extend(caption_refs)
+            _add_child_elements(pic, doc, obj, pelem)
+
+        elif ptype == "table":
+            current_list = None
+            text = ""
+            caption_refs = []
+            item_label = DocItemLabel(pelem["name"])
+
+            for caption in obj["captions"]:
+                text += caption["text"]
+
+                for nprov in caption["prov"]:
+                    npaths = nprov["$ref"].split("/")
+                    nelem = resolve_item(npaths, doc_glm)
+
+                    if nelem is None:
+                        # print(f"warning: undefined caption {npaths}")
+                        continue
+
+                    span_i = nelem["span"][0]
+                    span_j = nelem["span"][1]
+
+                    cap_text = caption["text"][span_i:span_j]
+
+                    # doc_glm["page-elements"].remove(nelem)
+
+                    prov = ProvenanceItem(
+                        page_no=nelem["page"],
+                        charspan=tuple(nelem["span"]),
+                        bbox=BoundingBox.from_tuple(
+                            nelem["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                        ),
+                    )
+
+                    caption_obj = doc.add_text(
+                        label=DocItemLabel.CAPTION, text=cap_text, prov=prov
+                    )
+                    caption_refs.append(caption_obj.get_ref())
+
+            table_cells_glm = _flatten_table_grid(obj["data"])
+
+            table_cells = []
+            for tbl_cell_glm in table_cells_glm:
+                if tbl_cell_glm["bbox"] is not None:
+                    bbox = BoundingBox.from_tuple(
+                        tbl_cell_glm["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                    )
+                else:
+                    bbox = None
+
+                is_col_header = False
+                is_row_header = False
+                is_row_section = False
+
+                if tbl_cell_glm["type"] == "col_header":
+                    is_col_header = True
+                elif tbl_cell_glm["type"] == "row_header":
+                    is_row_header = True
+                elif tbl_cell_glm["type"] == "row_section":
+                    is_row_section = True
+
+                table_cells.append(
+                    TableCell(
+                        row_span=tbl_cell_glm["row-span"][1]
+                        - tbl_cell_glm["row-span"][0],
+                        col_span=tbl_cell_glm["col-span"][1]
+                        - tbl_cell_glm["col-span"][0],
+                        start_row_offset_idx=tbl_cell_glm["row-span"][0],
+                        end_row_offset_idx=tbl_cell_glm["row-span"][1],
+                        start_col_offset_idx=tbl_cell_glm["col-span"][0],
+                        end_col_offset_idx=tbl_cell_glm["col-span"][1],
+                        text=tbl_cell_glm["text"],
+                        bbox=bbox,
+                        column_header=is_col_header,
+                        row_header=is_row_header,
+                        row_section=is_row_section,
+                    )
+                )
+
+            tbl_data = TableData(
+                num_rows=obj.get("#-rows", 0),
+                num_cols=obj.get("#-cols", 0),
+                table_cells=table_cells,
+            )
+
+            prov = ProvenanceItem(
+                page_no=pelem["page"],
+                charspan=(0, 0),
+                bbox=BoundingBox.from_tuple(
+                    pelem["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                ),
+            )
+
+            tbl = doc.add_table(data=tbl_data, prov=prov, label=item_label)
+            tbl.captions.extend(caption_refs)
+
+        elif ptype in [DocItemLabel.FORM.value, DocItemLabel.KEY_VALUE_REGION.value]:
+            label = DocItemLabel(ptype)
+            group_label = GroupLabel.UNSPECIFIED
+            if label == DocItemLabel.FORM:
+                group_label = GroupLabel.FORM_AREA
+            elif label == DocItemLabel.KEY_VALUE_REGION:
+                group_label = GroupLabel.KEY_VALUE_AREA
+
+            container_el = doc.add_group(label=group_label)
+
+            _add_child_elements(container_el, doc, obj, pelem)
+        elif "text" in obj:
+            text = obj["text"][span_i:span_j]
+
+            type_label = pelem["type"]
+            name_label = pelem["name"]
+            if update_name_label and len(props) > 0 and type_label == "paragraph":
+                prop = props[
+                    (props["type"] == "semantic") & (props["subj_path"] == iref)
+                ]
+                if len(prop) == 1 and prop.iloc[0]["confidence"] > 0.85:
+                    name_label = prop.iloc[0]["label"]
+
+            prov = ProvenanceItem(
+                page_no=pelem["page"],
+                charspan=(0, len(text)),
+                bbox=BoundingBox.from_tuple(
+                    pelem["bbox"], origin=CoordOrigin.BOTTOMLEFT
+                ),
+            )
+            label = DocItemLabel(name_label)
+
+            if label == DocItemLabel.LIST_ITEM:
+                if current_list is None:
+                    current_list = doc.add_group(label=GroupLabel.LIST, name="list")
+
+                # TODO: Infer if this is a numbered or a bullet list item
+                doc.add_list_item(
+                    text=text, enumerated=False, prov=prov, parent=current_list
+                )
+            elif label == DocItemLabel.SECTION_HEADER:
+                current_list = None
+
+                doc.add_heading(text=text, prov=prov)
+            elif label == DocItemLabel.CODE:
+                current_list = None
+
+                doc.add_code(text=text, prov=prov)
+            elif label == DocItemLabel.FORMULA:
+                current_list = None
+
+                doc.add_text(label=DocItemLabel.FORMULA, text="", orig=text, prov=prov)
+            elif label in [DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER]:
+                current_list = None
+
+                doc.add_text(
+                    label=DocItemLabel(name_label),
+                    text=text,
+                    prov=prov,
+                    content_layer=ContentLayer.FURNITURE,
+                )
+            else:
+                current_list = None
+
+                doc.add_text(label=DocItemLabel(name_label), text=text, prov=prov)
+
+    return doc
+
+
+def _add_child_elements(container_el, doc, obj, pelem):
+    payload = obj.get("payload")
+    if payload is not None:
+        children = payload.get("children", [])
+
+        for child in children:
+            c_label = DocItemLabel(child["label"])
+            c_bbox = BoundingBox.model_validate(child["bbox"]).to_bottom_left_origin(
+                doc.pages[pelem["page"]].size.height
+            )
+            c_text = " ".join(
+                [
+                    cell["text"].replace("\x02", "-").strip()
+                    for cell in child["cells"]
+                    if len(cell["text"].strip()) > 0
+                ]
+            )
+
+            c_prov = ProvenanceItem(
+                page_no=pelem["page"], charspan=(0, len(c_text)), bbox=c_bbox
+            )
+            if c_label == DocItemLabel.LIST_ITEM:
+                # TODO: Infer if this is a numbered or a bullet list item
+                doc.add_list_item(parent=container_el, text=c_text, prov=c_prov)
+            elif c_label == DocItemLabel.SECTION_HEADER:
+                doc.add_heading(parent=container_el, text=c_text, prov=c_prov)
+            else:
+                doc.add_text(
+                    parent=container_el, label=c_label, text=c_text, prov=c_prov
+                )