matcher.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. # copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # -*- encoding: utf-8 -*-
  15. import numpy as np
  16. from .matcher_utils import compute_iou, distance
  17. class TableMatch:
  18. def __init__(self, filter_ocr_result=True, use_master=False):
  19. self.filter_ocr_result = filter_ocr_result
  20. self.use_master = use_master
  21. def __call__(self, pred_structures, cell_bboxes, dt_boxes, rec_res):
  22. if self.filter_ocr_result:
  23. dt_boxes, rec_res = self._filter_ocr_result(cell_bboxes, dt_boxes, rec_res)
  24. matched_index = self.match_result(dt_boxes, cell_bboxes)
  25. pred_html, pred = self.get_pred_html(pred_structures, matched_index, rec_res)
  26. return pred_html
  27. def match_result(self, dt_boxes, cell_bboxes, min_iou=0.1**8):
  28. matched = {}
  29. for i, gt_box in enumerate(dt_boxes):
  30. distances = []
  31. for j, pred_box in enumerate(cell_bboxes):
  32. if len(pred_box) == 8:
  33. pred_box = [
  34. np.min(pred_box[0::2]),
  35. np.min(pred_box[1::2]),
  36. np.max(pred_box[0::2]),
  37. np.max(pred_box[1::2]),
  38. ]
  39. distances.append(
  40. (distance(gt_box, pred_box), 1.0 - compute_iou(gt_box, pred_box))
  41. ) # compute iou and l1 distance
  42. sorted_distances = distances.copy()
  43. # select det box by iou and l1 distance
  44. sorted_distances = sorted(
  45. sorted_distances, key=lambda item: (item[1], item[0])
  46. )
  47. # must > min_iou
  48. if sorted_distances[0][1] >= 1 - min_iou:
  49. continue
  50. if distances.index(sorted_distances[0]) not in matched:
  51. matched[distances.index(sorted_distances[0])] = [i]
  52. else:
  53. matched[distances.index(sorted_distances[0])].append(i)
  54. return matched
  55. def get_pred_html(self, pred_structures, matched_index, ocr_contents):
  56. end_html = []
  57. td_index = 0
  58. for tag in pred_structures:
  59. if "</td>" not in tag:
  60. end_html.append(tag)
  61. continue
  62. if "<td></td>" == tag:
  63. end_html.extend("<td>")
  64. if td_index in matched_index.keys():
  65. b_with = False
  66. if (
  67. "<b>" in ocr_contents[matched_index[td_index][0]]
  68. and len(matched_index[td_index]) > 1
  69. ):
  70. b_with = True
  71. end_html.extend("<b>")
  72. for i, td_index_index in enumerate(matched_index[td_index]):
  73. content = ocr_contents[td_index_index][0]
  74. if len(matched_index[td_index]) > 1:
  75. if len(content) == 0:
  76. continue
  77. if content[0] == " ":
  78. content = content[1:]
  79. if "<b>" in content:
  80. content = content[3:]
  81. if "</b>" in content:
  82. content = content[:-4]
  83. if len(content) == 0:
  84. continue
  85. if i != len(matched_index[td_index]) - 1 and " " != content[-1]:
  86. content += " "
  87. end_html.extend(content)
  88. if b_with:
  89. end_html.extend("</b>")
  90. if "<td></td>" == tag:
  91. end_html.append("</td>")
  92. else:
  93. end_html.append(tag)
  94. td_index += 1
  95. # Filter <thead></thead><tbody></tbody> elements
  96. filter_elements = ["<thead>", "</thead>", "<tbody>", "</tbody>"]
  97. end_html = [v for v in end_html if v not in filter_elements]
  98. return "".join(end_html), end_html
  99. def decode_logic_points(self, pred_structures):
  100. logic_points = []
  101. current_row = 0
  102. current_col = 0
  103. max_rows = 0
  104. max_cols = 0
  105. occupied_cells = {} # 用于记录已经被占用的单元格
  106. def is_occupied(row, col):
  107. return (row, col) in occupied_cells
  108. def mark_occupied(row, col, rowspan, colspan):
  109. for r in range(row, row + rowspan):
  110. for c in range(col, col + colspan):
  111. occupied_cells[(r, c)] = True
  112. i = 0
  113. while i < len(pred_structures):
  114. token = pred_structures[i]
  115. if token == "<tr>":
  116. current_col = 0 # 每次遇到 <tr> 时,重置当前列号
  117. elif token == "</tr>":
  118. current_row += 1 # 行结束,行号增加
  119. elif token.startswith("<td"):
  120. colspan = 1
  121. rowspan = 1
  122. j = i
  123. if token != "<td></td>":
  124. j += 1
  125. # 提取 colspan 和 rowspan 属性
  126. while j < len(pred_structures) and not pred_structures[
  127. j
  128. ].startswith(">"):
  129. if "colspan=" in pred_structures[j]:
  130. colspan = int(pred_structures[j].split("=")[1].strip("\"'"))
  131. elif "rowspan=" in pred_structures[j]:
  132. rowspan = int(pred_structures[j].split("=")[1].strip("\"'"))
  133. j += 1
  134. # 跳过已经处理过的属性 token
  135. i = j
  136. # 找到下一个未被占用的列
  137. while is_occupied(current_row, current_col):
  138. current_col += 1
  139. # 计算逻辑坐标
  140. r_start = current_row
  141. r_end = current_row + rowspan - 1
  142. col_start = current_col
  143. col_end = current_col + colspan - 1
  144. # 记录逻辑坐标
  145. logic_points.append([r_start, r_end, col_start, col_end])
  146. # 标记占用的单元格
  147. mark_occupied(r_start, col_start, rowspan, colspan)
  148. # 更新当前列号
  149. current_col += colspan
  150. # 更新最大行数和列数
  151. max_rows = max(max_rows, r_end + 1)
  152. max_cols = max(max_cols, col_end + 1)
  153. i += 1
  154. return logic_points
  155. def _filter_ocr_result(self, cell_bboxes, dt_boxes, rec_res):
  156. y1 = cell_bboxes[:, 1::2].min()
  157. new_dt_boxes = []
  158. new_rec_res = []
  159. for box, rec in zip(dt_boxes, rec_res):
  160. if np.max(box[1::2]) < y1:
  161. continue
  162. new_dt_boxes.append(box)
  163. new_rec_res.append(rec)
  164. return new_dt_boxes, new_rec_res