feat(reading-order): re-enable same-row left-to-right linking with additive r2l fix

docling_ibm_models/reading_order/reading_order_rb.py

@@ -320,12 +320,11 @@ def _init_l2r_map(
         state.l2r_map = {}
         state.r2l_map = {}
 
-        # this currently leads to errors ... might be necessary in the future ...
         for i, pelem_i in enumerate(page_elems):
             for j, pelem_j in enumerate(page_elems):
 
                 if (
-                    False  # pelem_i.follows_maintext_order(pelem_j)
+                    pelem_i.follows_maintext_order(pelem_j)
                     and pelem_i.is_strictly_left_of(pelem_j)
                     and pelem_i.overlaps_vertically_with_iou(pelem_j, 0.8)
                 ):
@@ -355,11 +354,12 @@ def _init_ud_maps(
 
         for j, pelem_j in enumerate(page_elems):
             if j in state.r2l_map:
-                i = state.r2l_map[j]
-                state.dn_map[i] = [j]
-                state.up_map[j] = [i]
-                continue
-
+                left_partner = state.r2l_map[j]
+                # Link the same-row left partner, then keep searching for vertical parents
+                if j not in state.dn_map[left_partner]:
+                    state.dn_map[left_partner].append(j)
+                if left_partner not in state.up_map[j]:
+                    state.up_map[j].append(left_partner)
             # Find elements above current that might precede it in reading order
             query_bbox = (pelem_j.l - 0.1, pelem_j.t, pelem_j.r + 0.1, float("inf"))
             candidates = list(spatial_idx.intersection(query_bbox))

tests/test_reading_order.py

@@ -247,6 +247,39 @@ def test_readingorder():
     print("  score(caption): ", mean_cp_score)
     print("score(footnotes): ", mean_ft_score)
 
+def test_caption_not_orphaned_in_two_column_figure():
+    """A caption above the right column of a same-row pair is read in place, not last."""
+    from docling_core.types.doc.base import CoordOrigin, Size
+    from docling_core.types.doc.labels import DocItemLabel
+
+    page_size = Size(width=600, height=800)
+
+    def elem(cid, label, l, r, b, t, text=""):
+        return PageElement(
+            cid=cid,
+            text=text,
+            page_no=1,
+            page_size=page_size,
+            label=label,
+            l=l,
+            r=r,
+            b=b,
+            t=t,
+            coord_origin=CoordOrigin.BOTTOMLEFT,
+        )
+    # picture (left) + caption (above the right column only); two same-row body columns below 
+    elements = [
+        elem(0, DocItemLabel.PICTURE, 60, 270, 400, 690),
+        elem(1, DocItemLabel.CAPTION, 340, 480, 300, 360, "Figure 1. Example"),
+        elem(2, DocItemLabel.TEXT, 60, 270, 200, 290, "left column body"),
+        elem(3, DocItemLabel.TEXT, 280, 494, 190, 290, "right column body text"),
+    ]
+
+    order = [e.cid for e in ReadingOrderPredictor().predict_reading_order(page_elements=elements)]
+
+    assert order.index(1) < order.index(3), (
+        f"caption (cid 1) should be read before the body columns, got {order}"
+    )
 
 """    
 def test_readingorder_multipage():
@@ -266,4 +299,4 @@ def test_readingorder_multipage():
     pred_elements = romodel.predict_reading_order(page_elements=true_elements)
     for true_elem, pred_elem in zip(true_elements, pred_elements):
         print("true: ", str(true_elem), ", pred: ", str(pred_elem))
-"""
+"""