Merge pull request #447 from glorialucheng/fix-bug

修复bug
netease-youdao · Jul 24, 2024 · deaaf43 · deaaf43
2 parents 8296f6d + 16555ba
commit deaaf43
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 7 deletions.
diff --git a/qanything_kernel/utils/loader/pdf_to_markdown/core/layout/table_rec/table_recover.py b/qanything_kernel/utils/loader/pdf_to_markdown/core/layout/table_rec/table_recover.py
@@ -30,6 +30,9 @@ def __call__(self, polygons: np.ndarray) -> Dict[int, Dict]:
     @staticmethod
     def get_rows(polygons: np.array) -> Dict[int, List[int]]:
         """对每个框进行行分类，框定哪个是一行的"""
+        if polygons.size == 0:
+            return {0: [0]}
+
         y_axis = polygons[:, 0, 1]
         if y_axis.size == 1:
             return {0: [0]}
@@ -43,6 +46,9 @@ def get_rows(polygons: np.array) -> Dict[int, List[int]]:
         if split_idxs.ndim == 0:
             split_idxs = split_idxs[None, ...]
 
+        if split_idxs.size == 0:
+            return {0: [0]}
+
         if max(split_idxs) != len(minus_res):
             split_idxs = np.append(split_idxs, len(minus_res))
 
@@ -58,6 +64,9 @@ def get_rows(polygons: np.array) -> Dict[int, List[int]]:
     def get_benchmark_cols(
         self, rows: Dict[int, List], polygons: np.ndarray
     ) -> Tuple[np.ndarray, List[float], int]:
+        if polygons.size == 0:
+            return None, [], 0
+
         longest_col = max(rows.values(), key=lambda x: len(x))
         longest_col_points = polygons[longest_col]
         longest_x = longest_col_points[:, 0, 0]
@@ -109,6 +118,9 @@ def get_benchmark_cols(
     def get_benchmark_rows(
         self, rows: Dict[int, List], polygons: np.ndarray
     ) -> Tuple[np.ndarray, List[float], int]:
+        if polygons.size == 0:
+            return None, []
+
         leftmost_cell_idxs = [v[0] for v in rows.values()]
         benchmark_x = polygons[leftmost_cell_idxs][:, 0, 1]
 
@@ -168,6 +180,9 @@ def get_merge_cells(
         each_col_widths: List[float],
         each_row_heights: List[float],
     ) -> Dict[int, Dict[int, int]]:
+        if polygons.size == 0:
+            return {}
+
         col_res_merge, row_res_merge = {}, {}
         merge_thresh = 20
         for cur_row, col_list in rows.items():

diff --git a/qanything_kernel/utils/loader/pdf_to_markdown/core/layout/table_rec/utils_table_recover.py b/qanything_kernel/utils/loader/pdf_to_markdown/core/layout/table_rec/utils_table_recover.py
@@ -125,6 +125,9 @@ def combine_two_poly(polygons: np.ndarray, idxs: np.ndarray) -> np.ndarray:
 def match_ocr_cell(
     polygons: np.ndarray, ocr_res: List[Union[List[List[float]], str, str]]
 ) -> Dict[int, List]:
+    if polygons.size == 0:
+        return {}, {}, {}
+
     cell_box_map = {}
     head_box_map = {}
     tail_box_map = {}
@@ -279,13 +282,16 @@ def plot_html_wireless_table(logi_points, cell_box_map):
 
     new_table_dict = {}
     for k, v in table_dict.items():
-        if len(col_dict[k]) < max(col_dict[k]):
-            for i in range(max(col_dict[k])+1):
-                if i not in col_dict[k]: 
-                    v.insert(i, '<td></td>')
-            new_table_dict[k] = ["<tr>"] + v + ["</tr>"]
-        else:
-            new_table_dict[k] = ["<tr>"] + v + ["</tr>"]
+        try:
+            if len(col_dict[k]) < max(col_dict[k]):
+                for i in range(max(col_dict[k])+1):
+                    if i not in col_dict[k]: 
+                        v.insert(i, '<td></td>')
+                new_table_dict[k] = ["<tr>"] + v + ["</tr>"]
+            else:
+                new_table_dict[k] = ["<tr>"] + v + ["</tr>"]
+        except Exception as e:
+            continue
 
     html_start = """<html><body><table><tbody>"""
     # html_start = """<html><style type="text/css">td {border-left: 1px solid;border-bottom:1px solid;}table, th {border-top:1px solid;font-size: 10px;border-collapse: collapse;border-right: 1px solid;}</style><body><table style="border-bottom:1px solid;border-top:1px solid;"><tbody>"""