diff --git a/docling_ibm_models/reading_order/reading_order_rb.py b/docling_ibm_models/reading_order/reading_order_rb.py index 88b150e..a57600a 100644 --- a/docling_ibm_models/reading_order/reading_order_rb.py +++ b/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)) diff --git a/tests/test_reading_order.py b/tests/test_reading_order.py index 568195d..3943bf4 100644 --- a/tests/test_reading_order.py +++ b/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)) -""" \ No newline at end of file +"""