table_recover.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  1. # -*- encoding: utf-8 -*-
  2. # @Author: SWHL
  3. # @Contact: liekkaskono@163.com
  4. from typing import Dict, List, Tuple
  5. import numpy as np
  6. class TableRecover:
  7. def __init__(
  8. self,
  9. ):
  10. pass
  11. def __call__(
  12. self, polygons: np.ndarray, rows_thresh=10, col_thresh=15
  13. ) -> Dict[int, Dict]:
  14. rows = self.get_rows(polygons, rows_thresh)
  15. longest_col, each_col_widths, col_nums = self.get_benchmark_cols(
  16. rows, polygons, col_thresh
  17. )
  18. each_row_heights, row_nums = self.get_benchmark_rows(rows, polygons)
  19. table_res, logic_points_dict = self.get_merge_cells(
  20. polygons,
  21. rows,
  22. row_nums,
  23. col_nums,
  24. longest_col,
  25. each_col_widths,
  26. each_row_heights,
  27. )
  28. logic_points = np.array(
  29. [logic_points_dict[i] for i in range(len(polygons))]
  30. ).astype(np.int32)
  31. return table_res, logic_points
  32. @staticmethod
  33. def get_rows(polygons: np.array, rows_thresh=10) -> Dict[int, List[int]]:
  34. """对每个框进行行分类,框定哪个是一行的"""
  35. y_axis = polygons[:, 0, 1]
  36. if y_axis.size == 1:
  37. return {0: [0]}
  38. concat_y = np.array(list(zip(y_axis, y_axis[1:])))
  39. minus_res = concat_y[:, 1] - concat_y[:, 0]
  40. result = {}
  41. split_idxs = np.argwhere(abs(minus_res) > rows_thresh).squeeze()
  42. # 如果都在一行,则将所有下标设置为同一行
  43. if split_idxs.size == 0:
  44. return {0: [i for i in range(len(y_axis))]}
  45. if split_idxs.ndim == 0:
  46. split_idxs = split_idxs[None, ...]
  47. if max(split_idxs) != len(minus_res):
  48. split_idxs = np.append(split_idxs, len(minus_res))
  49. start_idx = 0
  50. for row_num, idx in enumerate(split_idxs):
  51. if row_num != 0:
  52. start_idx = split_idxs[row_num - 1] + 1
  53. result.setdefault(row_num, []).extend(range(start_idx, idx + 1))
  54. # 计算每一行相邻cell的iou,如果大于0.2,则合并为同一个cell
  55. return result
  56. def get_benchmark_cols(
  57. self, rows: Dict[int, List], polygons: np.ndarray, col_thresh=15
  58. ) -> Tuple[np.ndarray, List[float], int]:
  59. longest_col = max(rows.values(), key=lambda x: len(x))
  60. longest_col_points = polygons[longest_col]
  61. longest_x_start = list(longest_col_points[:, 0, 0])
  62. longest_x_end = list(longest_col_points[:, 2, 0])
  63. min_x = longest_x_start[0]
  64. max_x = longest_x_end[-1]
  65. # 根据当前col的起始x坐标,更新col的边界
  66. # 2025.2.22 --- 解决最长列可能漏掉最后一列的问题
  67. def update_longest_col(col_x_list, cur_v, min_x_, max_x_, insert_last):
  68. for i, v in enumerate(col_x_list):
  69. if cur_v - col_thresh <= v <= cur_v + col_thresh:
  70. break
  71. if cur_v < min_x_:
  72. col_x_list.insert(0, cur_v)
  73. min_x_ = cur_v
  74. break
  75. if cur_v > max_x_:
  76. if insert_last:
  77. col_x_list.append(cur_v)
  78. max_x_ = cur_v
  79. break
  80. if cur_v < v:
  81. col_x_list.insert(i, cur_v)
  82. break
  83. return min_x_, max_x_
  84. for row_value in rows.values():
  85. cur_row_start = list(polygons[row_value][:, 0, 0])
  86. cur_row_end = list(polygons[row_value][:, 2, 0])
  87. for idx, (cur_v_start, cur_v_end) in enumerate(
  88. zip(cur_row_start, cur_row_end)
  89. ):
  90. min_x, max_x = update_longest_col(
  91. longest_x_start, cur_v_start, min_x, max_x, True
  92. )
  93. min_x, max_x = update_longest_col(
  94. longest_x_start, cur_v_end, min_x, max_x, False
  95. )
  96. longest_x_start = np.array(longest_x_start)
  97. each_col_widths = (longest_x_start[1:] - longest_x_start[:-1]).tolist()
  98. each_col_widths.append(max_x - longest_x_start[-1])
  99. col_nums = longest_x_start.shape[0]
  100. return longest_x_start, each_col_widths, col_nums
  101. def get_benchmark_rows(
  102. self, rows: Dict[int, List], polygons: np.ndarray
  103. ) -> Tuple[np.ndarray, List[float], int]:
  104. leftmost_cell_idxs = [v[0] for v in rows.values()]
  105. benchmark_x = polygons[leftmost_cell_idxs][:, 0, 1]
  106. each_row_widths = (benchmark_x[1:] - benchmark_x[:-1]).tolist()
  107. # 求出最后一行cell中,最大的高度作为最后一行的高度
  108. bottommost_idxs = list(rows.values())[-1]
  109. bottommost_boxes = polygons[bottommost_idxs]
  110. # fix self.compute_L2(v[3, :], v[0, :]), v为逆时针,即v[3]为右上,v[0]为左上,v[1]为左下
  111. max_height = max([self.compute_L2(v[1, :], v[0, :]) for v in bottommost_boxes])
  112. each_row_widths.append(max_height)
  113. row_nums = benchmark_x.shape[0]
  114. return each_row_widths, row_nums
  115. @staticmethod
  116. def compute_L2(a1: np.ndarray, a2: np.ndarray) -> float:
  117. return np.linalg.norm(a2 - a1)
  118. def get_merge_cells(
  119. self,
  120. polygons: np.ndarray,
  121. rows: Dict,
  122. row_nums: int,
  123. col_nums: int,
  124. longest_col: np.ndarray,
  125. each_col_widths: List[float],
  126. each_row_heights: List[float],
  127. ) -> Dict[int, Dict[int, int]]:
  128. col_res_merge, row_res_merge = {}, {}
  129. logic_points = {}
  130. merge_thresh = 10
  131. for cur_row, col_list in rows.items():
  132. one_col_result, one_row_result = {}, {}
  133. for one_col in col_list:
  134. box = polygons[one_col]
  135. box_width = self.compute_L2(box[3, :], box[0, :])
  136. # 不一定是从0开始的,应该综合已有值和x坐标位置来确定起始位置
  137. loc_col_idx = np.argmin(np.abs(longest_col - box[0, 0]))
  138. col_start = max(sum(one_col_result.values()), loc_col_idx)
  139. # 计算合并多少个列方向单元格
  140. for i in range(col_start, col_nums):
  141. col_cum_sum = sum(each_col_widths[col_start : i + 1])
  142. if i == col_start and col_cum_sum > box_width:
  143. one_col_result[one_col] = 1
  144. break
  145. elif abs(col_cum_sum - box_width) <= merge_thresh:
  146. one_col_result[one_col] = i + 1 - col_start
  147. break
  148. # 这里必须进行修正,不然会出现超越阈值范围后列交错
  149. elif col_cum_sum > box_width:
  150. idx = (
  151. i
  152. if abs(col_cum_sum - box_width)
  153. < abs(col_cum_sum - each_col_widths[i] - box_width)
  154. else i - 1
  155. )
  156. one_col_result[one_col] = idx + 1 - col_start
  157. break
  158. else:
  159. one_col_result[one_col] = col_nums - col_start
  160. col_end = one_col_result[one_col] + col_start - 1
  161. box_height = self.compute_L2(box[1, :], box[0, :])
  162. row_start = cur_row
  163. for j in range(row_start, row_nums):
  164. row_cum_sum = sum(each_row_heights[row_start : j + 1])
  165. # box_height 不确定是几行的高度,所以要逐个试验,找一个最近的几行的高
  166. # 如果第一次row_cum_sum就比box_height大,那么意味着?丢失了一行
  167. if j == row_start and row_cum_sum > box_height:
  168. one_row_result[one_col] = 1
  169. break
  170. elif abs(box_height - row_cum_sum) <= merge_thresh:
  171. one_row_result[one_col] = j + 1 - row_start
  172. break
  173. # 这里必须进行修正,不然会出现超越阈值范围后行交错
  174. elif row_cum_sum > box_height:
  175. idx = (
  176. j
  177. if abs(row_cum_sum - box_height)
  178. < abs(row_cum_sum - each_row_heights[j] - box_height)
  179. else j - 1
  180. )
  181. one_row_result[one_col] = idx + 1 - row_start
  182. break
  183. else:
  184. one_row_result[one_col] = row_nums - row_start
  185. row_end = one_row_result[one_col] + row_start - 1
  186. logic_points[one_col] = np.array(
  187. [row_start, row_end, col_start, col_end]
  188. )
  189. col_res_merge[cur_row] = one_col_result
  190. row_res_merge[cur_row] = one_row_result
  191. res = {}
  192. for i, (c, r) in enumerate(zip(col_res_merge.values(), row_res_merge.values())):
  193. res[i] = {k: [cc, r[k]] for k, cc in c.items()}
  194. return res, logic_points