feat: Introduce SegmentedPage for OCR (#91)

* Add README for Docling-DPBench

Signed-off-by: Christoph Auer <[email protected]>

* Establish SegmentedPage support in DatasetRecord and DatasetRecordWithPrediction

Signed-off-by: Christoph Auer <[email protected]>

* Add SegmentedPage usage to PixParse dataset provider

Signed-off-by: Christoph Auer <[email protected]>

* feat: Update FUNSDDatasetBuilder to support segmented pages (#93)

* feat: update FUNSDDatasetBuilder to support segmented pages

* address review comments

---------

Co-authored-by: samiullahchattha <[email protected]>

* feat: Update XFUNDDatasetBuilder to support segmented pages (#92)

* feat: update XFUNDDatasetBuilder to support segmented pages

* address review comments

---------

Co-authored-by: samiullahchattha <[email protected]>

---------

Signed-off-by: Christoph Auer <[email protected]>
Co-authored-by: samiuc <[email protected]>
Co-authored-by: samiullahchattha <[email protected]>

Author: 3 people

docling_eval/cli/main.py

@@ -254,6 +254,7 @@ def get_prediction_provider(
         pipeline_options.images_scale = 2.0
         pipeline_options.generate_page_images = True
         pipeline_options.generate_picture_images = True
+        pipeline_options.generate_parsed_pages = True
 
         if artifacts_path is not None:
             pipeline_options.artifacts_path = artifacts_path

docling_eval/datamodels/dataset_record.py

@@ -9,11 +9,14 @@
 from datasets import Sequence, Value
 from docling.datamodel.base_models import ConversionStatus
 from docling_core.types import DoclingDocument
+from docling_core.types.doc.page import SegmentedPage
 from docling_core.types.io import DocumentStream
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator
 
 from docling_eval.datamodels.types import EvaluationModality, PredictionFormats
 
+seg_adapter = TypeAdapter(Dict[int, SegmentedPage])
+
 
 class DatasetRecord(
     BaseModel
@@ -24,6 +27,9 @@ class DatasetRecord(
     doc_hash: Optional[str] = Field(alias="document_filehash", default=None)
 
     ground_truth_doc: DoclingDocument = Field(alias="GroundTruthDocument")
+    ground_truth_segmented_pages: Dict[int, SegmentedPage] = Field(
+        alias="ground_truth_segmented_pages", default={}
+    )
     original: Optional[Union[DocumentStream, Path]] = Field(
         alias="BinaryDocument", default=None
     )
@@ -53,6 +59,7 @@ def features(cls):
                 cls.get_field_alias("doc_path"): Value("string"),
                 cls.get_field_alias("doc_hash"): Value("string"),
                 cls.get_field_alias("ground_truth_doc"): Value("string"),
+                cls.get_field_alias("ground_truth_segmented_pages"): Value("string"),
                 cls.get_field_alias("ground_truth_pictures"): Sequence(
                     Features_Image()
                 ),
@@ -102,6 +109,9 @@ def as_record_dict(self):
                 self.ground_truth_doc.export_to_dict()
             ),
             self.get_field_alias("ground_truth_pictures"): self.ground_truth_pictures,
+            self.get_field_alias("ground_truth_segmented_pages"): seg_adapter.dump_json(
+                self.ground_truth_segmented_pages
+            ),
             self.get_field_alias(
                 "ground_truth_page_images"
             ): self.ground_truth_page_images,
@@ -143,6 +153,12 @@ def validate_record_dict(cls, data: dict):
         if gt_doc_alias in data and isinstance(data[gt_doc_alias], str):
             data[gt_doc_alias] = json.loads(data[gt_doc_alias])
 
+        gt_seg_pages_alias = cls.get_field_alias("ground_truth_segmented_pages")
+        if gt_seg_pages_alias in data and isinstance(data[gt_seg_pages_alias], bytes):
+            data[gt_seg_pages_alias] = seg_adapter.validate_json(
+                data[gt_seg_pages_alias]
+            )
+
         gt_page_img_alias = cls.get_field_alias("ground_truth_page_images")
         if gt_page_img_alias in data:
             for ix, item in enumerate(data[gt_page_img_alias]):
@@ -171,6 +187,11 @@ class DatasetRecordWithPrediction(DatasetRecord):
     predicted_doc: Optional[DoclingDocument] = Field(
         alias="PredictedDocument", default=None
     )
+
+    predicted_segmented_pages: Dict[int, SegmentedPage] = Field(
+        alias="predicted_segmented_pages", default={}
+    )
+
     original_prediction: Optional[str] = None
     prediction_format: PredictionFormats  # some enum type
     prediction_timings: Optional[Dict] = Field(alias="prediction_timings", default=None)
@@ -187,20 +208,22 @@ class DatasetRecordWithPrediction(DatasetRecord):
     @classmethod
     def features(cls):
         return {
-            cls.get_field_alias("predictor_info"): Value("string"),
-            cls.get_field_alias("status"): Value("string"),
             cls.get_field_alias("doc_id"): Value("string"),
             cls.get_field_alias("doc_path"): Value("string"),
             cls.get_field_alias("doc_hash"): Value("string"),
             cls.get_field_alias("ground_truth_doc"): Value("string"),
+            cls.get_field_alias("ground_truth_segmented_pages"): Value("string"),
             cls.get_field_alias("ground_truth_pictures"): Sequence(Features_Image()),
             cls.get_field_alias("ground_truth_page_images"): Sequence(Features_Image()),
-            cls.get_field_alias("predicted_doc"): Value("string"),
-            cls.get_field_alias("predicted_pictures"): Sequence(Features_Image()),
-            cls.get_field_alias("predicted_page_images"): Sequence(Features_Image()),
             cls.get_field_alias("original"): Value("string"),
             cls.get_field_alias("mime_type"): Value("string"),
             cls.get_field_alias("modalities"): Sequence(Value("string")),
+            cls.get_field_alias("predictor_info"): Value("string"),
+            cls.get_field_alias("status"): Value("string"),
+            cls.get_field_alias("predicted_doc"): Value("string"),
+            cls.get_field_alias("predicted_segmented_pages"): Value("string"),
+            cls.get_field_alias("predicted_pictures"): Sequence(Features_Image()),
+            cls.get_field_alias("predicted_page_images"): Sequence(Features_Image()),
             cls.get_field_alias("prediction_format"): Value("string"),
             cls.get_field_alias("prediction_timings"): Value("string"),
         }
@@ -211,6 +234,8 @@ def as_record_dict(self):
             {
                 self.get_field_alias("prediction_format"): self.prediction_format.value,
                 self.get_field_alias("prediction_timings"): self.prediction_timings,
+                self.get_field_alias("predictor_info"): self.predictor_info,
+                self.get_field_alias("status"): (self.status),
             }
         )
 
@@ -220,15 +245,16 @@ def as_record_dict(self):
                     self.get_field_alias("predicted_doc"): json.dumps(
                         self.predicted_doc.export_to_dict()
                     ),
+                    self.get_field_alias(
+                        "predicted_segmented_pages"
+                    ): seg_adapter.dump_json(self.predicted_segmented_pages),
                     self.get_field_alias("predicted_pictures"): self.predicted_pictures,
                     self.get_field_alias(
                         "predicted_page_images"
                     ): self.predicted_page_images,
                     self.get_field_alias("original_prediction"): (
                         self.original_prediction
                     ),
-                    self.get_field_alias("status"): (self.status),
-                    self.get_field_alias("predictor_info"): self.predictor_info,
                 }
             )
 
@@ -262,6 +288,14 @@ def validate_prediction_record_dict(cls, data: dict):
         if pred_doc_alias in data and isinstance(data[pred_doc_alias], str):
             data[pred_doc_alias] = json.loads(data[pred_doc_alias])
 
+        pred_seg_pages_alias = cls.get_field_alias("predicted_segmented_pages")
+        if pred_seg_pages_alias in data and isinstance(
+            data[pred_seg_pages_alias], bytes
+        ):
+            data[pred_seg_pages_alias] = seg_adapter.validate_json(
+                data[pred_seg_pages_alias]
+            )
+
         pred_page_img_alias = cls.get_field_alias("predicted_page_images")
         if pred_page_img_alias in data:
             for ix, item in enumerate(data[pred_page_img_alias]):

docling_eval/dataset_builders/funsd_builder.py

@@ -3,13 +3,20 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from datasets import DownloadManager
 from docling_core.types import DoclingDocument
 from docling_core.types.doc import BoundingBox, ImageRef, PageItem, ProvenanceItem, Size
 from docling_core.types.doc.document import GraphCell, GraphData, GraphLink
 from docling_core.types.doc.labels import GraphCellLabel, GraphLinkLabel
+from docling_core.types.doc.page import (
+    BoundingRectangle,
+    PageGeometry,
+    SegmentedPage,
+    TextCell,
+)
+from docling_core.types.io import DocumentStream
 from PIL import Image
 from tqdm import tqdm
 
@@ -193,9 +200,9 @@ def get_overall_bbox(
         bbox_instance = BoundingBox.enclosing_bbox(all_bboxes)
         return bbox_instance
 
-    def populate_key_value_item(
+    def _create_ground_truth_doc(
         self, doc: DoclingDocument, funsd_data: dict
-    ) -> DoclingDocument:
+    ) -> Tuple[DoclingDocument, Dict[int, SegmentedPage]]:
         """
         Populate the key-value item from the FUNSD data.
 
@@ -210,6 +217,19 @@ def populate_key_value_item(
             raise ValueError("Invalid FUNSD data: missing 'form' key.")
 
         form_items = funsd_data["form"]
+        segmented_pages: Dict[int, SegmentedPage] = {}
+
+        page_item: PageItem = doc.pages[1]
+        seg_page = SegmentedPage(
+            dimension=PageGeometry(
+                angle=0,
+                rect=BoundingRectangle.from_bounding_box(
+                    BoundingBox(
+                        l=0, t=0, r=page_item.size.width, b=page_item.size.height
+                    )
+                ),
+            )
+        )
 
         cell_by_id = {}
         for item in form_items:
@@ -242,6 +262,23 @@ def populate_key_value_item(
             )
             cell_by_id[cell_id] = cell
 
+            for word in item.get("words", []):
+                text = word.get("text", None)
+                bbox = word.get("box", None)
+                if bbox is None or text is None:
+                    continue
+                bbox_obj = self.convert_bbox(bbox)
+                seg_page.word_cells.append(
+                    TextCell(
+                        from_ocr=True,
+                        rect=BoundingRectangle.from_bounding_box(bbox_obj),
+                        text=text,
+                        orig=text,
+                    )
+                )
+
+        segmented_pages[doc.pages[1].page_no] = seg_page
+
         # unique linking pairs
         linking_set = set()
         for item in form_items:
@@ -283,7 +320,7 @@ def populate_key_value_item(
 
         sort_cell_ids(doc)
 
-        return doc
+        return doc, segmented_pages
 
     def iterate(self) -> Iterable[DatasetRecord]:
         """
@@ -358,23 +395,28 @@ def iterate(self) -> Iterable[DatasetRecord]:
                 true_doc.pages[1] = page_item
 
                 # Populate document with key-value data
-                true_doc = self.populate_key_value_item(true_doc, funsd_data)
+                true_doc, seg_pages = self._create_ground_truth_doc(
+                    true_doc, funsd_data
+                )
 
                 # Extract images
                 true_doc, true_pictures, true_page_images = extract_images(
                     document=true_doc,
                     pictures_column=BenchMarkColumns.GROUNDTRUTH_PICTURES.value,
                     page_images_column=BenchMarkColumns.GROUNDTRUTH_PAGE_IMAGES.value,
                 )
-
+                image_stream = DocumentStream(
+                    name=img_path.stem, stream=io.BytesIO(img_bytes)
+                )
                 # Create dataset record
                 record = DatasetRecord(
                     doc_id=img_path.stem,
                     doc_hash=get_binhash(img_bytes),
                     ground_truth_doc=true_doc,
-                    original=None,
+                    ground_truth_segmented_pages=seg_pages,
+                    original=image_stream,
                     mime_type="image/png",
-                    modalities=[EvaluationModality.KEY_VALUE],
+                    modalities=[EvaluationModality.KEY_VALUE, EvaluationModality.OCR],
                     ground_truth_pictures=true_pictures,
                     ground_truth_page_images=true_page_images,
                 )

docling_eval/dataset_builders/pixparse_builder.py

@@ -2,7 +2,7 @@
 import logging
 from io import BytesIO
 from pathlib import Path
-from typing import Dict, Iterable, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from docling_core.types import DoclingDocument
 from docling_core.types.doc import (
@@ -14,6 +14,12 @@
     ProvenanceItem,
     Size,
 )
+from docling_core.types.doc.page import (
+    BoundingRectangle,
+    PageGeometry,
+    SegmentedPage,
+    TextCell,
+)
 from docling_core.types.io import DocumentStream
 from PIL import Image
 from tqdm import tqdm
@@ -54,7 +60,7 @@ def __init__(
 
     def _create_ground_truth_doc(
         self, doc_id: str, gt_data: Dict, image: Image.Image
-    ) -> DoclingDocument:
+    ) -> Tuple[DoclingDocument, Dict[int, SegmentedPage]]:
         """Create a DoclingDocument from ground truth data and image file."""
         true_doc = DoclingDocument(name=doc_id)
 
@@ -72,8 +78,19 @@ def _create_ground_truth_doc(
         )
         true_doc.pages[1] = page_item
 
+        segmented_pages: Dict[int, SegmentedPage] = {}
+
         for page_idx, page in enumerate(gt_data["pages"], 1):
-            for text, bbox, _ in zip(page["text"], page["bbox"], page["score"]):
+            seg_page = SegmentedPage(
+                dimension=PageGeometry(
+                    angle=0,
+                    rect=BoundingRectangle.from_bounding_box(
+                        BoundingBox(l=0, t=0, r=image.width, b=image.height)
+                    ),
+                )
+            )
+
+            for text, bbox, score in zip(page["text"], page["bbox"], page["score"]):
                 bbox_obj = BoundingBox.from_tuple(
                     (
                         float(bbox[0]),
@@ -83,12 +100,18 @@ def _create_ground_truth_doc(
                     ),
                     CoordOrigin.TOPLEFT,
                 )
-                prov = ProvenanceItem(
-                    page_no=page_idx, bbox=bbox_obj, charspan=(0, len(text))
+                seg_page.textline_cells.append(
+                    TextCell(
+                        from_ocr=True,
+                        rect=BoundingRectangle.from_bounding_box(bbox_obj),
+                        text=text,
+                        orig=text,
+                        confidence=score,
+                    )
                 )
-                true_doc.add_text(label=DocItemLabel.TEXT, text=text, prov=prov)
+            segmented_pages[page_idx] = seg_page
 
-        return true_doc
+        return true_doc, segmented_pages
 
     def iterate(self) -> Iterable[DatasetRecord]:
         if not self.retrieved and self.must_retrieve:
@@ -135,7 +158,9 @@ def iterate(self) -> Iterable[DatasetRecord]:
                 ):
                     image = image.convert("RGB")
 
-                true_doc = self._create_ground_truth_doc(doc_id, gt_data, image)
+                true_doc, seg_pages = self._create_ground_truth_doc(
+                    doc_id, gt_data, image
+                )
 
                 # Extract images from the ground truth document
                 true_doc, true_pictures, true_page_images = extract_images(
@@ -158,6 +183,7 @@ def iterate(self) -> Iterable[DatasetRecord]:
                     doc_id=doc_id,
                     doc_hash=get_binhash(img_bytes),
                     ground_truth_doc=true_doc,
+                    ground_truth_segmented_pages=seg_pages,
                     original=image_stream,
                     mime_type="image/png",
                     modalities=[