table_recognition_post_processing_v2.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310
  1. # copyright (c) 2024 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. import math
  15. from typing import Any, Dict, Optional
  16. import numpy as np
  17. from ..layout_parsing.utils import get_sub_regions_ocr_res
  18. from ..components import convert_points_to_boxes
  19. from .result import SingleTableRecognitionResult
  20. from ..ocr.result import OCRResult
  21. def get_ori_image_coordinate(x: int, y: int, box_list: list) -> list:
  22. """
  23. get the original coordinate from Cropped image to Original image.
  24. Args:
  25. x (int): x coordinate of cropped image
  26. y (int): y coordinate of cropped image
  27. box_list (list): list of table bounding boxes, eg. [[x1, y1, x2, y2, x3, y3, x4, y4]]
  28. Returns:
  29. list: list of original coordinates, eg. [[x1, y1, x2, y2, x3, y3, x4, y4]]
  30. """
  31. if not box_list:
  32. return box_list
  33. offset = np.array([x, y] * 2)
  34. box_list = np.array(box_list)
  35. ori_box_list = offset + box_list
  36. return ori_box_list
  37. def convert_table_structure_pred_bbox(
  38. cell_points_list: list, crop_start_point: list, img_shape: tuple
  39. ) -> None:
  40. """
  41. Convert the predicted table structure bounding boxes to the original image coordinate system.
  42. Args:
  43. cell_points_list (list): Bounding boxes ('bbox').
  44. crop_start_point (list): A list of two integers representing the starting point (x, y) of the cropped image region.
  45. img_shape (tuple): A tuple of two integers representing the shape (height, width) of the original image.
  46. Returns:
  47. cell_points_list (list): Bounding boxes ('bbox').
  48. """
  49. ori_cell_points_list = get_ori_image_coordinate(
  50. crop_start_point[0], crop_start_point[1], cell_points_list
  51. )
  52. ori_cell_points_list = np.reshape(ori_cell_points_list, (-1, 4, 2))
  53. cell_box_list = convert_points_to_boxes(ori_cell_points_list)
  54. img_height, img_width = img_shape
  55. cell_box_list = np.clip(
  56. cell_box_list, 0, [img_width, img_height, img_width, img_height]
  57. )
  58. return cell_box_list
  59. def distance(box_1: list, box_2: list) -> float:
  60. """
  61. compute the distance between two boxes
  62. Args:
  63. box_1 (list): first rectangle box,eg.(x1, y1, x2, y2)
  64. box_2 (list): second rectangle box,eg.(x1, y1, x2, y2)
  65. Returns:
  66. float: the distance between two boxes
  67. """
  68. x1, y1, x2, y2 = box_1
  69. x3, y3, x4, y4 = box_2
  70. center1_x = (x1 + x2) / 2
  71. center1_y = (y1 + y2) / 2
  72. center2_x = (x3 + x4) / 2
  73. center2_y = (y3 + y4) / 2
  74. dis = math.sqrt((center2_x - center1_x) ** 2 + (center2_y - center1_y) ** 2)
  75. dis_2 = abs(x3 - x1) + abs(y3 - y1)
  76. dis_3 = abs(x4 - x2) + abs(y4 - y2)
  77. return dis + min(dis_2, dis_3)
  78. def compute_iou(rec1: list, rec2: list) -> float:
  79. """
  80. computing IoU
  81. Args:
  82. rec1 (list): (x1, y1, x2, y2)
  83. rec2 (list): (x1, y1, x2, y2)
  84. Returns:
  85. float: Intersection over Union
  86. """
  87. # computing area of each rectangles
  88. S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1])
  89. S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1])
  90. # computing the sum_area
  91. sum_area = S_rec1 + S_rec2
  92. # find the each edge of intersect rectangle
  93. left_line = max(rec1[0], rec2[0])
  94. right_line = min(rec1[2], rec2[2])
  95. top_line = max(rec1[1], rec2[1])
  96. bottom_line = min(rec1[3], rec2[3])
  97. # judge if there is an intersect
  98. if left_line >= right_line or top_line >= bottom_line:
  99. return 0.0
  100. else:
  101. intersect = (right_line - left_line) * (bottom_line - top_line)
  102. return (intersect / (sum_area - intersect)) * 1.0
  103. def match_table_and_ocr(cell_box_list: list, ocr_dt_boxes: list) -> dict:
  104. """
  105. match table and ocr
  106. Args:
  107. cell_box_list (list): bbox for table cell, 2 points, [left, top, right, bottom]
  108. ocr_dt_boxes (list): bbox for ocr, 2 points, [left, top, right, bottom]
  109. Returns:
  110. dict: matched dict, key is table index, value is ocr index
  111. """
  112. matched = {}
  113. for i, ocr_box in enumerate(np.array(ocr_dt_boxes)):
  114. ocr_box = ocr_box.astype(np.float32)
  115. distances = []
  116. for j, table_box in enumerate(cell_box_list):
  117. if len(table_box) == 8:
  118. table_box = [
  119. np.min(table_box[0::2]),
  120. np.min(table_box[1::2]),
  121. np.max(table_box[0::2]),
  122. np.max(table_box[1::2]),
  123. ]
  124. distances.append(
  125. (distance(table_box, ocr_box), 1.0 - compute_iou(table_box, ocr_box))
  126. ) # compute iou and l1 distance
  127. sorted_distances = distances.copy()
  128. # select det box by iou and l1 distance
  129. sorted_distances = sorted(sorted_distances, key=lambda item: (item[1], item[0]))
  130. if distances.index(sorted_distances[0]) not in matched.keys():
  131. matched[distances.index(sorted_distances[0])] = [i]
  132. else:
  133. matched[distances.index(sorted_distances[0])].append(i)
  134. return matched
  135. def get_html_result(
  136. matched_index: dict, ocr_contents: dict, pred_structures: list
  137. ) -> str:
  138. """
  139. Generates HTML content based on the matched index, OCR contents, and predicted structures.
  140. Args:
  141. matched_index (dict): A dictionary containing matched indices.
  142. ocr_contents (dict): A dictionary of OCR contents.
  143. pred_structures (list): A list of predicted HTML structures.
  144. Returns:
  145. str: Generated HTML content as a string.
  146. """
  147. pred_html = []
  148. td_index = 0
  149. head_structure = pred_structures[0:3]
  150. html = "".join(head_structure)
  151. table_structure = pred_structures[3:-3]
  152. for tag in table_structure:
  153. if "</td>" in tag:
  154. if "<td></td>" == tag:
  155. pred_html.extend("<td>")
  156. if td_index in matched_index.keys():
  157. b_with = False
  158. if (
  159. "<b>" in ocr_contents[matched_index[td_index][0]]
  160. and len(matched_index[td_index]) > 1
  161. ):
  162. b_with = True
  163. pred_html.extend("<b>")
  164. for i, td_index_index in enumerate(matched_index[td_index]):
  165. content = ocr_contents[td_index_index]
  166. if len(matched_index[td_index]) > 1:
  167. if len(content) == 0:
  168. continue
  169. if content[0] == " ":
  170. content = content[1:]
  171. if "<b>" in content:
  172. content = content[3:]
  173. if "</b>" in content:
  174. content = content[:-4]
  175. if len(content) == 0:
  176. continue
  177. if i != len(matched_index[td_index]) - 1 and " " != content[-1]:
  178. content += " "
  179. pred_html.extend(content)
  180. if b_with:
  181. pred_html.extend("</b>")
  182. if "<td></td>" == tag:
  183. pred_html.append("</td>")
  184. else:
  185. pred_html.append(tag)
  186. td_index += 1
  187. else:
  188. pred_html.append(tag)
  189. html += "".join(pred_html)
  190. end_structure = pred_structures[-3:]
  191. html += "".join(end_structure)
  192. return html
  193. def sort_table_cells_boxes(boxes):
  194. """
  195. Sort the input list of bounding boxes by using the DBSCAN algorithm to cluster based on the top-left y-coordinate (y1), and then sort within each line from left to right based on the top-left x-coordinate (x1).
  196. Args:
  197. boxes (list of lists): The input list of bounding boxes, where each bounding box is formatted as [x1, y1, x2, y2].
  198. Returns:
  199. sorted_boxes (list of lists): The list of bounding boxes sorted.
  200. """
  201. import numpy as np
  202. from sklearn.cluster import DBSCAN
  203. # Extract the top-left y-coordinates (y1)
  204. y1_coords = np.array([box[1] for box in boxes])
  205. y1_coords = y1_coords.reshape(-1, 1) # Convert to a 2D array
  206. # Choose an appropriate eps parameter based on the range of y-values
  207. y_range = y1_coords.max() - y1_coords.min()
  208. eps = y_range / 50 # Adjust the denominator as needed
  209. # Perform clustering using DBSCAN
  210. db = DBSCAN(eps=eps, min_samples=1).fit(y1_coords)
  211. labels = db.labels_
  212. # Group bounding boxes by their labels
  213. clusters = {}
  214. for label, box in zip(labels, boxes):
  215. if label not in clusters:
  216. clusters[label] = []
  217. clusters[label].append(box)
  218. # Sort rows based on y-coordinates
  219. # Compute the average y1 value for each row and sort from top to bottom
  220. sorted_rows = sorted(
  221. clusters.items(), key=lambda item: np.mean([box[1] for box in item[1]])
  222. )
  223. # Within each row, sort by x1 coordinate
  224. sorted_boxes = []
  225. for label, row in sorted_rows:
  226. row_sorted = sorted(row, key=lambda x: x[0])
  227. sorted_boxes.extend(row_sorted)
  228. return sorted_boxes
  229. def get_table_recognition_res(
  230. table_box: list,
  231. table_structure_result: list,
  232. table_cells_result: list,
  233. overall_ocr_res: OCRResult,
  234. ) -> SingleTableRecognitionResult:
  235. """
  236. Retrieve table recognition result from cropped image info, table structure prediction, and overall OCR result.
  237. Args:
  238. table_box (list): Information about the location of cropped image, including the bounding box.
  239. table_structure_result (list): Predicted table structure.
  240. table_cells_result (list): Predicted table cells.
  241. overall_ocr_res (OCRResult): Overall OCR result from the input image.
  242. Returns:
  243. SingleTableRecognitionResult: An object containing the single table recognition result.
  244. """
  245. table_box = np.array([table_box])
  246. table_ocr_pred = get_sub_regions_ocr_res(overall_ocr_res, table_box)
  247. crop_start_point = [table_box[0][0], table_box[0][1]]
  248. img_shape = overall_ocr_res["doc_preprocessor_res"]["output_img"].shape[0:2]
  249. ori_table_cells = convert_table_structure_pred_bbox(
  250. table_cells_result, crop_start_point, img_shape
  251. )
  252. ocr_dt_boxes = table_ocr_pred["rec_boxes"]
  253. ocr_texts_res = table_ocr_pred["rec_texts"]
  254. table_cells_result = sort_table_cells_boxes(table_cells_result)
  255. ocr_dt_boxes = sort_table_cells_boxes(ocr_dt_boxes)
  256. matched_index = match_table_and_ocr(table_cells_result, ocr_dt_boxes)
  257. pred_html = get_html_result(matched_index, ocr_texts_res, table_structure_result)
  258. single_img_res = {
  259. "cell_box_list": ori_table_cells,
  260. "table_ocr_pred": table_ocr_pred,
  261. "pred_html": pred_html,
  262. }
  263. return SingleTableRecognitionResult(single_img_res)