fix: chain merge continuations instead of stopping after one pair (#156)

Signed-off-by: stone <frank.schruefer@t-online.de>

Author: Frank-Schruefer

docling_ibm_models/reading_order/reading_order_rb.py

@@ -174,6 +174,15 @@ def predict_merges(
 
         merges: Dict[int, List[int]] = {}
 
+        skip_labels = [
+            DocItemLabel.PAGE_HEADER,
+            DocItemLabel.PAGE_FOOTER,
+            DocItemLabel.TABLE,
+            DocItemLabel.PICTURE,
+            DocItemLabel.CAPTION,
+            DocItemLabel.FOOTNOTE,
+        ]
+
         curr_ind = -1
         for ind, elem in enumerate(sorted_elements):
 
@@ -182,32 +191,43 @@ def predict_merges(
 
             if elem.label in [DocItemLabel.TEXT]:
 
-                ind_p1 = ind + 1
-                while ind_p1 < len(sorted_elements) and sorted_elements[ind_p1] in [
-                    DocItemLabel.PAGE_HEADER,
-                    DocItemLabel.PAGE_FOOTER,
-                    DocItemLabel.TABLE,
-                    DocItemLabel.PICTURE,
-                    DocItemLabel.CAPTION,
-                    DocItemLabel.FOOTNOTE,
-                ]:
-                    ind_p1 += 1
-
-                if (
-                    ind_p1 < len(sorted_elements)
-                    and sorted_elements[ind_p1].label == elem.label
-                    and (
-                        elem.page_no != sorted_elements[ind_p1].label
-                        or elem.is_strictly_left_of(sorted_elements[ind_p1])
-                    )
-                ):
-
-                    m1 = re.fullmatch(r".+([a-z,\-])(\s*)", elem.text)
-                    m2 = re.fullmatch(r"(\s*[a-z])(.+)", sorted_elements[ind_p1].text)
+                merge_list: List[int] = []
+                check_ind = ind
+
+                while True:
+                    ind_p1 = check_ind + 1
+                    while (
+                        ind_p1 < len(sorted_elements)
+                        and sorted_elements[ind_p1].label in skip_labels
+                    ):
+                        ind_p1 += 1
+
+                    if (
+                        ind_p1 < len(sorted_elements)
+                        and sorted_elements[ind_p1].label == elem.label
+                        and (
+                            elem.page_no != sorted_elements[ind_p1].page_no
+                            or elem.is_strictly_left_of(sorted_elements[ind_p1])
+                        )
+                    ):
+                        m1 = re.fullmatch(
+                            r".+([a-z,\-])(\s*)", sorted_elements[check_ind].text
+                        )
+                        m2 = re.fullmatch(
+                            r"(\s*[a-z])(.+)", sorted_elements[ind_p1].text
+                        )
+
+                        if m1 and m2:
+                            merge_list.append(sorted_elements[ind_p1].cid)
+                            curr_ind = ind_p1
+                            check_ind = ind_p1
+                        else:
+                            break
+                    else:
+                        break
 
-                    if m1 and m2:
-                        merges[elem.cid] = [sorted_elements[ind_p1].cid]
-                        curr_ind = ind_p1
+                if merge_list:
+                    merges[elem.cid] = merge_list
 
         return merges
 
@@ -681,7 +701,7 @@ def _find_to_captions(
         """
 
         def _remove_overlapping_indexes(
-            mapping: Dict[int, List[int]]
+            mapping: Dict[int, List[int]],
         ) -> Dict[int, List[int]]:
             used = set()
             result = {}

docling_ibm_models/tableformer/data_management/tf_predictor.py

@@ -433,7 +433,7 @@ def resize_img(self, image, width=None, height=None, inter=cv2.INTER_AREA):
         # initialize the dimensions of the image to be resized and
         # grab the image size
         dim = None
-        (h, w) = image.shape[:2]
+        h, w = image.shape[:2]
         sf = 1.0
         # if both the width and height are None, then return the
         # original image