From d00109737653a7e670322100540e1dc06adf015a Mon Sep 17 00:00:00 2001 From: Cesar Berrospi Ramis <75900930+ceberam@users.noreply.github.com> Date: Wed, 9 Apr 2025 17:40:40 +0200 Subject: [PATCH] docling(xlsx): add bounding boxes and page size information in cell units Signed-off-by: Cesar Berrospi Ramis <75900930+ceberam@users.noreply.github.com> --- docling/backend/msexcel_backend.py | 88 +++++++++++++++---- .../groundtruth/docling_v2/test-01.xlsx.json | 68 +++++++------- tests/test_backend_msexcel.py | 19 +++- 3 files changed, 122 insertions(+), 53 deletions(-) diff --git a/docling/backend/msexcel_backend.py b/docling/backend/msexcel_backend.py index 722bcb90..971b93cd 100644 --- a/docling/backend/msexcel_backend.py +++ b/docling/backend/msexcel_backend.py @@ -1,11 +1,12 @@ import logging from io import BytesIO from pathlib import Path -from typing import Any, Union +from typing import Any, Union, cast from docling_core.types.doc import ( BoundingBox, CoordOrigin, + DocItem, DoclingDocument, DocumentOrigin, GroupLabel, @@ -16,9 +17,11 @@ from docling_core.types.doc import ( TableData, ) from openpyxl import load_workbook +from openpyxl.drawing.image import Image +from openpyxl.drawing.spreadsheet_drawing import TwoCellAnchor from openpyxl.worksheet.worksheet import Worksheet from PIL import Image as PILImage -from pydantic import BaseModel +from pydantic import BaseModel, NonNegativeInt, PositiveInt from typing_extensions import override from docling.backend.abstract_backend import ( @@ -50,21 +53,37 @@ class ExcelCell(BaseModel): class ExcelTable(BaseModel): - """Represents an Excel table. + """Represents an Excel table on a worksheet. Attributes: + anchor: The column and row indices of the upper-left cell of the table + (0-based index). num_rows: The number of rows in the table. num_cols: The number of columns in the table. data: The data in the table, represented as a list of ExcelCell objects. """ + anchor: tuple[NonNegativeInt, NonNegativeInt] num_rows: int num_cols: int data: list[ExcelCell] class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBackend): - """Backend for parsing Excel workbooks.""" + """Backend for parsing Excel workbooks. + + The backend converts an Excel workbook into a DoclingDocument object. + Each worksheet is converted into a separate page. + The following elements are parsed: + - Cell contents, parsed as tables. If two groups of cells are disconnected + between each other, they will be parsed as two different tables. + - Images, parsed as PictureItem objects. + + The DoclingDocument tables and pictures have their provenance information, including + the position in their original Excel worksheet. The position is represented by a + bounding box object with the cell indices as units (0-based index). The size of this + bounding box is the number of columns and rows that the table or picture spans. + """ @override def __init__( @@ -106,7 +125,7 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken @override def is_valid(self) -> bool: - _log.info(f"valid: {self.valid}") + _log.debug(f"valid: {self.valid}") return self.valid @classmethod @@ -170,19 +189,19 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken for sheet_name in self.workbook.sheetnames: _log.info(f"Processing sheet: {sheet_name}") - # Access the sheet by name sheet = self.workbook[sheet_name] - idx = self.workbook.index(sheet) - # TODO: check concept of Size as number of rows and cols - doc.add_page(page_no=idx + 1, size=Size()) + page_no = self.workbook.index(sheet) + 1 + # do not rely on sheet.max_column, sheet.max_row if there are images + page = doc.add_page(page_no=page_no, size=Size(width=0, height=0)) self.parents[0] = doc.add_group( parent=None, label=GroupLabel.SECTION, name=f"sheet: {sheet_name}", ) - doc = self._convert_sheet(doc, sheet) + width, height = self._find_page_size(doc, page_no) + page.size = Size(width=width, height=height) else: _log.error("Workbook is not initialized.") @@ -222,6 +241,8 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken tables = self._find_data_tables(sheet) for excel_table in tables: + origin_col = excel_table.anchor[0] + origin_row = excel_table.anchor[1] num_rows = excel_table.num_rows num_cols = excel_table.num_cols @@ -254,7 +275,13 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken page_no=page_no, charspan=(0, 0), bbox=BoundingBox.from_tuple( - (0, 0, 0, 0), origin=CoordOrigin.BOTTOMLEFT + ( + origin_col, + origin_row, + origin_col + num_cols, + origin_row + num_rows, + ), + origin=CoordOrigin.TOPLEFT, ), ), ) @@ -322,7 +349,6 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken row_span = 1 col_span = 1 - # _log.info(sheet.merged_cells.ranges) for merged_range in sheet.merged_cells.ranges: if ( @@ -346,7 +372,6 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken col_span=col_span, ) ) - # _log.info(f"cell: {ri}, {rj} -> {ri - start_row}, {rj - start_col}, {row_span}, {col_span}: {str(cell.value)}") # Mark all cells in the span as visited for span_row in range(ri, ri + row_span): @@ -355,6 +380,7 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken return ( ExcelTable( + anchor=(start_col, start_row), num_rows=max_row + 1 - start_row, num_cols=max_col + 1 - start_col, data=data, @@ -446,13 +472,21 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken Returns: The updated DoclingDocument. """ - if self.workbook is not None: # Iterate over byte images in the sheet - for image in sheet._images: # type: ignore[attr-defined] + for item in sheet._images: # type: ignore[attr-defined] try: - pil_image = PILImage.open(image.ref) + image: Image = cast(Image, item) + pil_image = PILImage.open(image.ref) # type: ignore[arg-type] page_no = self.workbook.index(sheet) + 1 + anchor = (0, 0, 0, 0) + if isinstance(image.anchor, TwoCellAnchor): + anchor = ( + image.anchor._from.col, + image.anchor._from.row, + image.anchor.to.col + 1, + image.anchor.to.row + 1, + ) doc.add_picture( parent=self.parents[0], image=ImageRef.from_pil(image=pil_image, dpi=72), @@ -461,7 +495,7 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken page_no=page_no, charspan=(0, 0), bbox=BoundingBox.from_tuple( - (0, 0, 0, 0), origin=CoordOrigin.TOPLEFT + anchor, origin=CoordOrigin.TOPLEFT ), ), ) @@ -469,3 +503,23 @@ class MsExcelDocumentBackend(DeclarativeDocumentBackend, PaginatedDocumentBacken _log.error("could not extract the image from excel sheets") return doc + + @staticmethod + def _find_page_size( + doc: DoclingDocument, page_no: PositiveInt + ) -> tuple[float, float]: + left: float = -1.0 + top: float = -1.0 + right: float = -1.0 + bottom: float = -1.0 + for item, _ in doc.iterate_items(traverse_pictures=True, page_no=page_no): + if not isinstance(item, DocItem): + continue + for provenance in item.prov: + bbox = provenance.bbox + left = min(left, bbox.l) if left != -1 else bbox.l + right = max(right, bbox.r) if right != -1 else bbox.r + top = min(top, bbox.t) if top != -1 else bbox.t + bottom = max(bottom, bbox.b) if bottom != -1 else bbox.b + + return (right - left, bottom - top) diff --git a/tests/data/groundtruth/docling_v2/test-01.xlsx.json b/tests/data/groundtruth/docling_v2/test-01.xlsx.json index bf6c19c6..2bdfe509 100644 --- a/tests/data/groundtruth/docling_v2/test-01.xlsx.json +++ b/tests/data/groundtruth/docling_v2/test-01.xlsx.json @@ -101,10 +101,10 @@ { "page_no": 3, "bbox": { - "l": 0.0, - "t": 0.0, - "r": 0.0, - "b": 0.0, + "l": 8.0, + "t": 18.0, + "r": 13.0, + "b": 36.0, "coord_origin": "TOPLEFT" }, "charspan": [ @@ -143,9 +143,9 @@ "bbox": { "l": 0.0, "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "r": 3.0, + "b": 7.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -697,9 +697,9 @@ "bbox": { "l": 0.0, "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "r": 4.0, + "b": 9.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -1613,11 +1613,11 @@ { "page_no": 2, "bbox": { - "l": 0.0, - "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "l": 6.0, + "t": 4.0, + "r": 9.0, + "b": 9.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -2019,11 +2019,11 @@ { "page_no": 2, "bbox": { - "l": 0.0, - "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "l": 2.0, + "t": 13.0, + "r": 5.0, + "b": 18.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -2427,9 +2427,9 @@ "bbox": { "l": 0.0, "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "r": 3.0, + "b": 7.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -2907,11 +2907,11 @@ { "page_no": 3, "bbox": { - "l": 0.0, - "t": 0.0, - "r": 0.0, - "b": 0.0, - "coord_origin": "BOTTOMLEFT" + "l": 4.0, + "t": 6.0, + "r": 7.0, + "b": 13.0, + "coord_origin": "TOPLEFT" }, "charspan": [ 0, @@ -3383,22 +3383,22 @@ "pages": { "1": { "size": { - "width": 0.0, - "height": 0.0 + "width": 3.0, + "height": 7.0 }, "page_no": 1 }, "2": { "size": { - "width": 0.0, - "height": 0.0 + "width": 9.0, + "height": 18.0 }, "page_no": 2 }, "3": { "size": { - "width": 0.0, - "height": 0.0 + "width": 13.0, + "height": 36.0 }, "page_no": 3 } diff --git a/tests/test_backend_msexcel.py b/tests/test_backend_msexcel.py index 86b0c013..0604429c 100644 --- a/tests/test_backend_msexcel.py +++ b/tests/test_backend_msexcel.py @@ -57,7 +57,7 @@ def documents() -> list[tuple[Path, DoclingDocument]]: return documents -def test_e2e_xlsx_conversions(documents): +def test_e2e_xlsx_conversions(documents) -> None: for gt_path, doc in documents: pred_md: str = doc.export_to_markdown() assert verify_export(pred_md, str(gt_path) + ".md"), "export to md" @@ -74,7 +74,13 @@ def test_e2e_xlsx_conversions(documents): ), "document document" -def test_page_count(): +def test_pages(documents) -> None: + """Test the page count and page size of converted documents. + + Args: + documents: The paths and converted documents. + """ + # number of pages from the backend method path = [item for item in get_xlsx_paths() if item.stem == "test-01"][0] in_doc = InputDocument( path_or_stream=path, @@ -84,3 +90,12 @@ def test_page_count(): ) backend = MsExcelDocumentBackend(in_doc=in_doc, path_or_stream=path) assert backend.page_count() == 3 + + # number of pages from the converted document + doc = [item for path, item in documents if path.stem == "test-01"][0] + assert len(doc.pages) == 3 + + # page sizes as number of cells + assert doc.pages.get(1).size.as_tuple() == (3.0, 7.0) + assert doc.pages.get(2).size.as_tuple() == (9.0, 18.0) + assert doc.pages.get(3).size.as_tuple() == (13.0, 36.0)