diff --git a/docling_core/transforms/visualizer/table_visualizer.py b/docling_core/transforms/visualizer/table_visualizer.py index 4e3adf20..f0cc6774 100644 --- a/docling_core/transforms/visualizer/table_visualizer.py +++ b/docling_core/transforms/visualizer/table_visualizer.py @@ -117,6 +117,7 @@ def _draw_table_rows( page_height: float, scale_x: float, scale_y: float, + doc: Optional[DoclingDocument] = None, ): """Draw individual table rows.""" # Create transparent overlay for proper alpha compositing @@ -124,7 +125,7 @@ def _draw_table_rows( overlay_draw = ImageDraw.Draw(overlay) main_draw = ImageDraw.Draw(page_image) - rows = table.data.get_row_bounding_boxes(minimal=self.params.minimal_row_bboxes) + rows = table.data.get_row_bounding_boxes(minimal=self.params.minimal_row_bboxes, doc=doc) for rid, bbox in rows.items(): tl_bbox = bbox.to_top_left_origin(page_height=page_height) @@ -152,6 +153,7 @@ def _draw_table_cols( page_height: float, scale_x: float, scale_y: float, + doc: Optional[DoclingDocument] = None, ): """Draw individual table columns.""" # Create transparent overlay for proper alpha compositing @@ -159,7 +161,7 @@ def _draw_table_cols( overlay_draw = ImageDraw.Draw(overlay) main_draw = ImageDraw.Draw(page_image) - cols = table.data.get_column_bounding_boxes(minimal=self.params.minimal_col_bboxes) + cols = table.data.get_column_bounding_boxes(minimal=self.params.minimal_col_bboxes, doc=doc) for cid, bbox in cols.items(): tl_bbox = bbox.to_top_left_origin(page_height=page_height) @@ -233,6 +235,7 @@ def _draw_doc_tables( page_image=image, scale_x=image.width / doc.pages[page_nr].size.width, scale_y=image.height / doc.pages[page_nr].size.height, + doc=doc, ) if self.params.show_cols: @@ -242,6 +245,7 @@ def _draw_doc_tables( page_image=image, scale_x=image.width / doc.pages[page_nr].size.width, scale_y=image.height / doc.pages[page_nr].size.height, + doc=doc, ) else: diff --git a/docling_core/types/doc/document.py b/docling_core/types/doc/document.py index 98a659f0..1c2fe871 100644 --- a/docling_core/types/doc/document.py +++ b/docling_core/types/doc/document.py @@ -369,6 +369,9 @@ def from_dict_format(cls, data: Any) -> Any: def _get_text(self, doc: Optional["DoclingDocument"] = None, **kwargs: Any) -> str: return self.text + def _get_bbox(self, doc: Optional["DoclingDocument"] = None, **kwargs: Any) -> Optional[BoundingBox]: + return self.bbox + class RichTableCell(TableCell): """RichTableCell.""" @@ -386,6 +389,37 @@ def _get_text(self, doc: Optional["DoclingDocument"] = None, **kwargs: Any) -> s else: return "" + @override + def _get_bbox(self, doc: Optional["DoclingDocument"] = None, **kwargs: Any) -> Optional[BoundingBox]: + # An explicit cell bbox always wins. + if self.bbox is not None: + return self.bbox + + # Otherwise derive the bbox from the referenced item: take the union of + # the provenance bounding boxes of the referenced item and all of its + # descendants. + if doc is None: + return None + + item = self.ref.resolve(doc=doc) + if not isinstance(item, NodeItem): + return None + + bboxes: list[BoundingBox] = [] + for node, _ in doc.iterate_items( + root=item, + with_groups=True, + traverse_pictures=True, + included_content_layers=set(ContentLayer), + ): + if isinstance(node, DocItem): + bboxes.extend(prov.bbox for prov in node.prov) + + if not bboxes: + return None + + return BoundingBox.enclosing_bbox(bboxes) + AnyTableCell = Annotated[ Union[RichTableCell, TableCell], @@ -598,7 +632,9 @@ def add_row(self, row: list[str]) -> None: """ self.insert_row(row_index=self.num_rows - 1, row=row, after=True) - def get_row_bounding_boxes(self, *, minimal: bool = True) -> dict[int, BoundingBox]: + def get_row_bounding_boxes( + self, *, minimal: bool = True, doc: Optional["DoclingDocument"] = None + ) -> dict[int, BoundingBox]: """Get the bounding box for each row in the table. Layout follows the table's ``orientation`` field: ``ROT_0`` / ``ROT_180`` @@ -612,16 +648,18 @@ def get_row_bounding_boxes(self, *, minimal: bool = True) -> dict[int, BoundingB row based on its cells. If False, all rows will have a uniform extent perpendicular to the row direction (l/r for ROT_0/ROT_180, t/b for ROT_90/ROT_270). + doc: Optional document used to resolve the bounding box of cells that + do not carry an explicit ``bbox`` (e.g. ``RichTableCell``), by + taking the union of the referenced item's provenance boxes. Returns: dict[int, BoundingBox]: A dictionary mapping row indices to their bounding boxes. Only rows with cells that have bounding boxes are included. """ horizontal = self.orientation in (Orientation.ROT_0, Orientation.ROT_180) - coords = [] - for cell in self.table_cells: - if cell.bbox is not None: - coords.append(cell.bbox.coord_origin) + # Resolve each cell's bbox once (RichTableCell may need to traverse `doc`). + cell_bboxes: list[Optional[BoundingBox]] = [cell._get_bbox(doc=doc) for cell in self.table_cells] + coords = [bbox.coord_origin for bbox in cell_bboxes if bbox is not None] if len(set(coords)) > 1: raise ValueError( @@ -635,13 +673,13 @@ def get_row_bounding_boxes(self, *, minimal: bool = True) -> dict[int, BoundingB row_cells_with_bbox: dict[int, list[BoundingBox]] = {} # Collect all cells in this row that have bounding boxes - for cell in self.table_cells: - if cell.bbox is not None and cell.start_row_offset_idx <= row_idx < cell.end_row_offset_idx: + for cell, cell_bbox in zip(self.table_cells, cell_bboxes): + if cell_bbox is not None and cell.start_row_offset_idx <= row_idx < cell.end_row_offset_idx: row_span = cell.end_row_offset_idx - cell.start_row_offset_idx if row_span in row_cells_with_bbox: - row_cells_with_bbox[row_span].append(cell.bbox) + row_cells_with_bbox[row_span].append(cell_bbox) else: - row_cells_with_bbox[row_span] = [cell.bbox] + row_cells_with_bbox[row_span] = [cell_bbox] # Calculate the enclosing bounding box for this row if len(row_cells_with_bbox) > 0: @@ -690,7 +728,9 @@ def get_row_bounding_boxes(self, *, minimal: bool = True) -> dict[int, BoundingB return row_bboxes - def get_column_bounding_boxes(self, *, minimal: bool = True) -> dict[int, BoundingBox]: + def get_column_bounding_boxes( + self, *, minimal: bool = True, doc: Optional["DoclingDocument"] = None + ) -> dict[int, BoundingBox]: """Get the bounding box for each column in the table. Layout follows the table's ``orientation`` field: ``ROT_0`` / ``ROT_180`` @@ -704,16 +744,18 @@ def get_column_bounding_boxes(self, *, minimal: bool = True) -> dict[int, Boundi column based on its cells. If False, all columns will have a uniform extent perpendicular to the column direction (t/b for ROT_0/ROT_180, l/r for ROT_90/ROT_270). + doc: Optional document used to resolve the bounding box of cells that + do not carry an explicit ``bbox`` (e.g. ``RichTableCell``), by + taking the union of the referenced item's provenance boxes. Returns: dict[int, BoundingBox]: A dictionary mapping column indices to their bounding boxes. Only columns with cells that have bounding boxes are included. """ horizontal = self.orientation in (Orientation.ROT_0, Orientation.ROT_180) - coords = [] - for cell in self.table_cells: - if cell.bbox is not None: - coords.append(cell.bbox.coord_origin) + # Resolve each cell's bbox once (RichTableCell may need to traverse `doc`). + cell_bboxes: list[Optional[BoundingBox]] = [cell._get_bbox(doc=doc) for cell in self.table_cells] + coords = [bbox.coord_origin for bbox in cell_bboxes if bbox is not None] if len(set(coords)) > 1: raise ValueError( @@ -727,13 +769,13 @@ def get_column_bounding_boxes(self, *, minimal: bool = True) -> dict[int, Boundi col_cells_with_bbox: dict[int, list[BoundingBox]] = {} # Collect all cells in this row that have bounding boxes - for cell in self.table_cells: - if cell.bbox is not None and cell.start_col_offset_idx <= col_idx < cell.end_col_offset_idx: + for cell, cell_bbox in zip(self.table_cells, cell_bboxes): + if cell_bbox is not None and cell.start_col_offset_idx <= col_idx < cell.end_col_offset_idx: col_span = cell.end_col_offset_idx - cell.start_col_offset_idx if col_span in col_cells_with_bbox: - col_cells_with_bbox[col_span].append(cell.bbox) + col_cells_with_bbox[col_span].append(cell_bbox) else: - col_cells_with_bbox[col_span] = [cell.bbox] + col_cells_with_bbox[col_span] = [cell_bbox] # Calculate the enclosing bounding box for this row if len(col_cells_with_bbox) > 0: diff --git a/test/test_docling_doc.py b/test/test_docling_doc.py index 3d723d76..1095a4b1 100644 --- a/test/test_docling_doc.py +++ b/test/test_docling_doc.py @@ -2658,6 +2658,85 @@ def test_table_data_horizontal_bounding_boxes_with_spans_no_overlap(): assert a.intersection_area_with(b) == 0, f"col overlap (minimal={minimal})" +def test_table_data_bounding_boxes_with_rich_table_cell(): + """RichTableCell without an own bbox contributes its referenced extent. + + A ``RichTableCell`` has an optional ``bbox``. When it is ``None`` but its + ``ref`` resolves to an item whose descendants carry provenance bboxes, the + row/column boxes must enclose that referenced extent -- but only when a + ``doc`` is supplied so the ref can be resolved. + """ + doc = DoclingDocument(name="rich-table") + doc.add_page(page_no=1, size=Size(width=200.0, height=200.0)) + + # Group with two text items carrying prov bboxes; the rich cell refs the group. + grp = doc.add_group(name="cellgrp") + t1 = doc.add_text( + label=DocItemLabel.TEXT, + text="a", + parent=grp, + prov=ProvenanceItem( + page_no=1, + bbox=BoundingBox(l=10, t=190, r=40, b=170, coord_origin=CoordOrigin.BOTTOMLEFT), + charspan=(0, 1), + ), + ) + doc.add_text( + label=DocItemLabel.TEXT, + text="b", + parent=grp, + prov=ProvenanceItem( + page_no=1, + bbox=BoundingBox(l=50, t=180, r=90, b=160, coord_origin=CoordOrigin.BOTTOMLEFT), + charspan=(0, 1), + ), + ) + + rich = RichTableCell( + ref=grp.get_ref(), + text="", + bbox=None, + start_row_offset_idx=0, + end_row_offset_idx=1, + start_col_offset_idx=0, + end_col_offset_idx=1, + ) + plain = TableCell( + text="x", + bbox=BoundingBox(l=100, t=190, r=150, b=160, coord_origin=CoordOrigin.BOTTOMLEFT), + start_row_offset_idx=0, + end_row_offset_idx=1, + start_col_offset_idx=1, + end_col_offset_idx=2, + ) + table = TableData(num_rows=1, num_cols=2, table_cells=[rich, plain]) + + # _get_bbox: no doc -> own (None) bbox; with doc -> union of referenced prov. + assert rich._get_bbox() is None + assert plain._get_bbox() == plain.bbox + rich_bbox = rich._get_bbox(doc=doc) + assert rich_bbox is not None + assert (rich_bbox.l, rich_bbox.r, rich_bbox.t, rich_bbox.b) == (10, 90, 190, 160) + + # Without doc the rich cell is skipped: row 0 only covers the plain cell. + rows_no_doc = table.get_row_bounding_boxes() + assert (rows_no_doc[0].l, rows_no_doc[0].r) == (100, 150) + # Column 0 (rich-only) has no resolvable bbox without a doc. + assert 0 not in table.get_column_bounding_boxes() + + # With doc the rich cell's referenced extent is included. + rows = table.get_row_bounding_boxes(doc=doc) + assert (rows[0].l, rows[0].r) == (10, 150) + cols = table.get_column_bounding_boxes(doc=doc) + assert (cols[0].l, cols[0].r) == (10, 90) + assert (cols[1].l, cols[1].r) == (100, 150) + + # An explicit bbox on the rich cell wins over the referenced extent. + rich.bbox = BoundingBox(l=5, t=195, r=95, b=155, coord_origin=CoordOrigin.BOTTOMLEFT) + assert rich._get_bbox(doc=doc) == rich.bbox + assert t1.parent is not None # sanity: ref tree wired up + + def _validate_doc(doc: DoclingDocument) -> DoclingDocument: return DoclingDocument.model_validate(doc.model_dump(mode="json"))