unet_table.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. import html
  2. import logging
  3. import os
  4. import time
  5. import traceback
  6. from dataclasses import dataclass, asdict
  7. from pathlib import Path
  8. from typing import List, Optional, Union, Dict, Any
  9. import cv2
  10. import numpy as np
  11. from loguru import logger
  12. from mineru.utils.enum_class import ModelPath
  13. from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
  14. from .table_structure_unet import TSRUnet
  15. from .table_recover import TableRecover
  16. from .wired_table_rec_utils import InputType, LoadImage
  17. from .table_recover_utils import (
  18. match_ocr_cell,
  19. plot_html_table,
  20. box_4_2_poly_to_box_4_1,
  21. sorted_ocr_boxes,
  22. gather_ocr_list_by_row,
  23. )
  24. @dataclass
  25. class UnetTableInput:
  26. model_path: str
  27. device: str = "cpu"
  28. @dataclass
  29. class UnetTableOutput:
  30. pred_html: Optional[str] = None
  31. cell_bboxes: Optional[np.ndarray] = None
  32. logic_points: Optional[np.ndarray] = None
  33. elapse: Optional[float] = None
  34. class UnetTableRecognition:
  35. def __init__(self, config: UnetTableInput):
  36. self.table_structure = TSRUnet(asdict(config))
  37. self.load_img = LoadImage()
  38. self.table_recover = TableRecover()
  39. def __call__(
  40. self,
  41. img: InputType,
  42. ocr_result: Optional[List[Union[List[List[float]], str, str]]] = None,
  43. **kwargs,
  44. ) -> UnetTableOutput:
  45. s = time.perf_counter()
  46. need_ocr = True
  47. col_threshold = 15
  48. row_threshold = 10
  49. if kwargs:
  50. need_ocr = kwargs.get("need_ocr", True)
  51. col_threshold = kwargs.get("col_threshold", 15)
  52. row_threshold = kwargs.get("row_threshold", 10)
  53. img = self.load_img(img)
  54. polygons, rotated_polygons = self.table_structure(img, **kwargs)
  55. if polygons is None:
  56. logging.warning("polygons is None.")
  57. return UnetTableOutput("", None, None, 0.0)
  58. try:
  59. table_res, logi_points = self.table_recover(
  60. rotated_polygons, row_threshold, col_threshold
  61. )
  62. # 将坐标由逆时针转为顺时针方向,后续处理与无线表格对齐
  63. polygons[:, 1, :], polygons[:, 3, :] = (
  64. polygons[:, 3, :].copy(),
  65. polygons[:, 1, :].copy(),
  66. )
  67. if not need_ocr:
  68. sorted_polygons, idx_list = sorted_ocr_boxes(
  69. [box_4_2_poly_to_box_4_1(box) for box in polygons]
  70. )
  71. return UnetTableOutput(
  72. "",
  73. sorted_polygons,
  74. logi_points[idx_list],
  75. time.perf_counter() - s,
  76. )
  77. cell_box_det_map, not_match_orc_boxes = match_ocr_cell(ocr_result, polygons)
  78. # 如果有识别框没有ocr结果,直接进行rec补充
  79. cell_box_det_map = self.fill_blank_rec(img, polygons, cell_box_det_map)
  80. # 转换为中间格式,修正识别框坐标,将物理识别框,逻辑识别框,ocr识别框整合为dict,方便后续处理
  81. t_rec_ocr_list = self.transform_res(cell_box_det_map, polygons, logi_points)
  82. # 将每个单元格中的ocr识别结果排序和同行合并,输出的html能完整保留文字的换行格式
  83. t_rec_ocr_list = self.sort_and_gather_ocr_res(t_rec_ocr_list)
  84. # cell_box_map =
  85. logi_points = [t_box_ocr["t_logic_box"] for t_box_ocr in t_rec_ocr_list]
  86. cell_box_det_map = {
  87. i: [ocr_box_and_text[1] for ocr_box_and_text in t_box_ocr["t_ocr_res"]]
  88. for i, t_box_ocr in enumerate(t_rec_ocr_list)
  89. }
  90. pred_html = plot_html_table(logi_points, cell_box_det_map)
  91. polygons = np.array(polygons).reshape(-1, 8)
  92. logi_points = np.array(logi_points)
  93. elapse = time.perf_counter() - s
  94. except Exception:
  95. logging.warning(traceback.format_exc())
  96. return UnetTableOutput("", None, None, 0.0)
  97. return UnetTableOutput(pred_html, polygons, logi_points, elapse)
  98. def transform_res(
  99. self,
  100. cell_box_det_map: Dict[int, List[any]],
  101. polygons: np.ndarray,
  102. logi_points: List[np.ndarray],
  103. ) -> List[Dict[str, any]]:
  104. res = []
  105. for i in range(len(polygons)):
  106. ocr_res_list = cell_box_det_map.get(i)
  107. if not ocr_res_list:
  108. continue
  109. xmin = min([ocr_box[0][0][0] for ocr_box in ocr_res_list])
  110. ymin = min([ocr_box[0][0][1] for ocr_box in ocr_res_list])
  111. xmax = max([ocr_box[0][2][0] for ocr_box in ocr_res_list])
  112. ymax = max([ocr_box[0][2][1] for ocr_box in ocr_res_list])
  113. dict_res = {
  114. # xmin,xmax,ymin,ymax
  115. "t_box": [xmin, ymin, xmax, ymax],
  116. # row_start,row_end,col_start,col_end
  117. "t_logic_box": logi_points[i].tolist(),
  118. # [[xmin,xmax,ymin,ymax], text]
  119. "t_ocr_res": [
  120. [box_4_2_poly_to_box_4_1(ocr_det[0]), ocr_det[1]]
  121. for ocr_det in ocr_res_list
  122. ],
  123. }
  124. res.append(dict_res)
  125. return res
  126. def sort_and_gather_ocr_res(self, res):
  127. for i, dict_res in enumerate(res):
  128. _, sorted_idx = sorted_ocr_boxes(
  129. [ocr_det[0] for ocr_det in dict_res["t_ocr_res"]], threshold=0.3
  130. )
  131. dict_res["t_ocr_res"] = [dict_res["t_ocr_res"][i] for i in sorted_idx]
  132. dict_res["t_ocr_res"] = gather_ocr_list_by_row(
  133. dict_res["t_ocr_res"], threshold=0.3
  134. )
  135. return res
  136. def fill_blank_rec(
  137. self,
  138. img: np.ndarray,
  139. sorted_polygons: np.ndarray,
  140. cell_box_map: Dict[int, List[str]],
  141. ) -> Dict[int, List[Any]]:
  142. """找到poly对应为空的框,尝试将直接将poly框直接送到识别中"""
  143. for i in range(sorted_polygons.shape[0]):
  144. if cell_box_map.get(i):
  145. continue
  146. box = sorted_polygons[i]
  147. cell_box_map[i] = [[box, "", 1]]
  148. continue
  149. return cell_box_map
  150. def escape_html(input_string):
  151. """Escape HTML Entities."""
  152. return html.escape(input_string)
  153. class UnetTableModel:
  154. def __init__(self, ocr_engine):
  155. model_path = os.path.join(auto_download_and_get_model_root_path(ModelPath.unet_structure), ModelPath.unet_structure)
  156. input_args = UnetTableInput(model_path=model_path)
  157. self.table_model = UnetTableRecognition(input_args)
  158. self.ocr_engine = ocr_engine
  159. def predict(self, img):
  160. bgr_img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
  161. ocr_result = self.ocr_engine.ocr(bgr_img)[0]
  162. if ocr_result:
  163. ocr_result = [
  164. [item[0], escape_html(item[1][0]), item[1][1]]
  165. for item in ocr_result
  166. if len(item) == 2 and isinstance(item[1], tuple)
  167. ]
  168. else:
  169. ocr_result = None
  170. if ocr_result:
  171. try:
  172. table_results = self.table_model(np.asarray(img), ocr_result)
  173. html_code = table_results.pred_html
  174. table_cell_bboxes = table_results.cell_bboxes
  175. logic_points = table_results.logic_points
  176. elapse = table_results.elapse
  177. return html_code, table_cell_bboxes, logic_points, elapse
  178. except Exception as e:
  179. logger.exception(e)
  180. return None, None, None, None