fix: add bottom bbox padding to prevent model from dropping last table row

Author: AnkitAhlawat7742

docling/models/stages/table_structure/table_structure_model.py

@@ -200,14 +200,22 @@ def predict_tables(
                 table_prediction = TableStructurePrediction()
                 page.predictions.tablestructure = table_prediction
 
+                # Small bottom padding (page units) added to the table bbox #3402
+                _BBOX_BOTTOM_PAD = 5.0
                 in_tables = [
                     (
                         cluster,
                         [
                             round(cluster.bbox.l) * self.scale,
                             round(cluster.bbox.t) * self.scale,
                             round(cluster.bbox.r) * self.scale,
-                            round(cluster.bbox.b) * self.scale,
+                            round(
+                                min(
+                                    cluster.bbox.b + _BBOX_BOTTOM_PAD,
+                                    page.size.height,
+                                )
+                            )
+                            * self.scale,
                         ],
                     )
                     for cluster in page.predictions.layout.clusters
@@ -227,7 +235,6 @@ def predict_tables(
                 for table_cluster, tbl_box in in_tables:
                     # Check if word-level cells are available from backend:
                     sp = page._backend.get_segmented_page()
-                    using_word_cells = False
                     if sp is not None:
                         tcells = sp.get_cells_in_bbox(
                             cell_unit=TextCellUnit.WORD,
@@ -236,8 +243,6 @@ def predict_tables(
                         if len(tcells) == 0:
                             # In case word-level cells yield empty
                             tcells = table_cluster.cells
-                        else:
-                            using_word_cells = True
                     else:
                         # Otherwise - we use normal (line/phrase) cells
                         tcells = table_cluster.cells
@@ -287,77 +292,6 @@ def predict_tables(
                         .get("rs_seq", [])
                     )
 
-                    # docling-side workaround for the edge case, main root cause is in TFPredictor
-                    # model doesn't see the last row because it's at the very bottom edge of the crop
-                    if (
-                        self.do_cell_matching
-                        and using_word_cells
-                        and tokens
-                        and table_cells
-                    ):
-                        matched_ids = {
-                            str(k)
-                            for k in table_out["predict_details"]
-                            .get("matches", {})
-                            .keys()
-                        }
-                        cells_bottom = max(
-                            (tc.bbox.b for tc in table_cells if tc.bbox is not None),
-                            default=0.0,
-                        )
-                        unmatched = []
-                        if cells_bottom > 0.0:
-                            for tok in tokens:
-                                if str(tok["id"]) in matched_ids:
-                                    continue
-                                tok_bbox = BoundingBox.model_validate(
-                                    tok["bbox"]
-                                ).scaled(1 / self.scale)
-                                if (
-                                    tok_bbox.t > cells_bottom
-                                    and tok_bbox.b <= table_cluster.bbox.b
-                                ):
-                                    unmatched.append((tok, tok_bbox))
-                        if len(unmatched) > 1:
-                            centres = sorted((tb.t + tb.b) / 2.0 for _, tb in unmatched)
-                            median_centre = centres[len(centres) // 2]
-                            avg_height = sum(tb.b - tb.t for _, tb in unmatched) / len(
-                                unmatched
-                            )
-                            coherent = all(
-                                abs((tb.t + tb.b) / 2.0 - median_centre) <= avg_height
-                                for _, tb in unmatched
-                            )
-                            if not coherent:
-                                unmatched = []
-                        if unmatched:
-                            extra_row_idx = num_rows
-                            num_cols = max(num_cols, 1)
-                            for col_idx, (tok, tok_bbox) in enumerate(unmatched):
-                                table_cells.append(
-                                    TableCell(
-                                        text=tok["text"],
-                                        bbox=tok_bbox,
-                                        row_span=1,
-                                        col_span=1,
-                                        start_row_offset_idx=extra_row_idx,
-                                        end_row_offset_idx=extra_row_idx + 1,
-                                        start_col_offset_idx=col_idx % num_cols,
-                                        end_col_offset_idx=(col_idx % num_cols) + 1,
-                                        column_header=False,
-                                        row_header=False,
-                                        row_section=False,
-                                    )
-                                )
-                            num_rows += 1
-                            _log.debug(
-                                "Recovered %d unmatched token(s) into extra row %d "
-                                "for table cluster id=%d",
-                                len(unmatched),
-                                extra_row_idx,
-                                table_cluster.id,
-                            )
-
                     tbl = Table(
                         otsl_seq=otsl_seq,
                         table_cells=table_cells,

tests/data/groundtruth/docling_v2/2203.01017v2.doctags.txt

@@ -9,7 +9,7 @@
 <section_header_level_1><loc_258><loc_138><loc_334><loc_143>a. Picture of a table:</section_header_level_1>
 <otsl><loc_258><loc_146><loc_439><loc_191><ched>1<nl></otsl>
 <text><loc_437><loc_224><loc_441><loc_228>7</text>
-<otsl><loc_258><loc_274><loc_439><loc_313><fcel>0<fcel>1 2 1<lcel><lcel><nl><fcel>3 4<fcel>5<fcel>6<fcel>7<nl><fcel>9 13<fcel>10<fcel>11<fcel>12<nl><fcel>8 2<fcel>14<fcel>15<fcel>16<nl><fcel>17<fcel>18<fcel>19<fcel>20<nl></otsl>
+<otsl><loc_258><loc_274><loc_439><loc_313><fcel>0<fcel>1<fcel>2 1<lcel><lcel><nl><fcel>3<fcel>4<fcel>5<fcel>6<fcel>7<nl><fcel>8 2<fcel>9<fcel>10<fcel>11<fcel>12<nl><ucel><fcel>13<fcel>14<fcel>15<fcel>16<nl><ucel><fcel>17<fcel>18<fcel>19<fcel>20<nl></otsl>
 <picture><loc_257><loc_143><loc_439><loc_313><caption><loc_252><loc_325><loc_445><loc_353>Figure 1: Picture of a table with subtle, complex features such as (1) multi-column headers, (2) cell with multi-row text and (3) cells with no content. Image from PubTabNet evaluation set, filename: 'PMC2944238 004 02'.</caption></picture>
 <text><loc_252><loc_369><loc_445><loc_420>Recently, significant progress has been made with vision based approaches to extract tables in documents. For the sake of completeness, the issue of table extraction from documents is typically decomposed into two separate challenges, i.e. (1) finding the location of the table(s) on a document-page and (2) finding the structure of a given table in the document.</text>
 <text><loc_252><loc_422><loc_445><loc_450><loc_41><loc_48><loc_234><loc_61>The first problem is called table-location and has been previously addressed [30, 38, 19, 21, 23, 26, 8] with stateof-the-art object-detection networks (e.g. YOLO and later on Mask-RCNN [9]). For all practical purposes, it can be considered as a solved problem, given enough ground-truth data to train on.</text>