table_structure_unet.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. import copy
  2. import math
  3. from typing import Optional, Dict, Any, Tuple
  4. import cv2
  5. import numpy as np
  6. from skimage import measure
  7. from mineru.model.utils.onnx_config import get_op_num_threads
  8. from .utils import OrtInferSession, resize_img
  9. from .utils_table_line_rec import (
  10. get_table_line,
  11. final_adjust_lines,
  12. min_area_rect_box,
  13. draw_lines,
  14. adjust_lines,
  15. )
  16. from.utils_table_recover import (
  17. sorted_ocr_boxes,
  18. box_4_2_poly_to_box_4_1,
  19. )
  20. class TSRUnet:
  21. def __init__(self, config: Dict):
  22. self.K = 1000
  23. self.MK = 4000
  24. self.mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
  25. self.std = np.array([58.395, 57.12, 57.375], dtype=np.float32)
  26. self.inp_height = 1024
  27. self.inp_width = 1024
  28. config["intra_op_num_threads"] = get_op_num_threads("MINERU_INTRA_OP_NUM_THREADS")
  29. config["inter_op_num_threads"] = get_op_num_threads("MINERU_INTER_OP_NUM_THREADS")
  30. self.session = OrtInferSession(config)
  31. def __call__(
  32. self, img: np.ndarray, **kwargs
  33. ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
  34. img_info = self.preprocess(img)
  35. pred = self.infer(img_info)
  36. polygons, rotated_polygons = self.postprocess(img, pred, **kwargs)
  37. if polygons.size == 0:
  38. return None, None
  39. polygons = polygons.reshape(polygons.shape[0], 4, 2)
  40. polygons[:, 3, :], polygons[:, 1, :] = (
  41. polygons[:, 1, :].copy(),
  42. polygons[:, 3, :].copy(),
  43. )
  44. rotated_polygons = rotated_polygons.reshape(rotated_polygons.shape[0], 4, 2)
  45. rotated_polygons[:, 3, :], rotated_polygons[:, 1, :] = (
  46. rotated_polygons[:, 1, :].copy(),
  47. rotated_polygons[:, 3, :].copy(),
  48. )
  49. _, idx = sorted_ocr_boxes(
  50. [box_4_2_poly_to_box_4_1(poly_box) for poly_box in rotated_polygons],
  51. threhold=0.4,
  52. )
  53. polygons = polygons[idx]
  54. rotated_polygons = rotated_polygons[idx]
  55. return polygons, rotated_polygons
  56. def preprocess(self, img) -> Dict[str, Any]:
  57. scale = (self.inp_height, self.inp_width)
  58. img, _, _ = resize_img(img, scale, True)
  59. img = img.copy().astype(np.float32)
  60. assert img.dtype != np.uint8
  61. mean = np.float64(self.mean.reshape(1, -1))
  62. stdinv = 1 / np.float64(self.std.reshape(1, -1))
  63. cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) # inplace
  64. cv2.subtract(img, mean, img) # inplace
  65. cv2.multiply(img, stdinv, img) # inplace
  66. img = img.transpose(2, 0, 1)
  67. images = img[None, :]
  68. return {"img": images}
  69. def infer(self, input):
  70. result = self.session(input["img"][None, ...])[0][0]
  71. result = result[0].astype(np.uint8)
  72. return result
  73. def postprocess(self, img, pred, **kwargs):
  74. row = kwargs.get("row", 50) if kwargs else 50
  75. col = kwargs.get("col", 30) if kwargs else 30
  76. h_lines_threshold = kwargs.get("h_lines_threshold", 100) if kwargs else 100
  77. v_lines_threshold = kwargs.get("v_lines_threshold", 15) if kwargs else 15
  78. angle = kwargs.get("angle", 50) if kwargs else 50
  79. enhance_box_line = kwargs.get("enhance_box_line", True) if kwargs else True
  80. morph_close = (
  81. kwargs.get("morph_close", enhance_box_line) if kwargs else enhance_box_line
  82. ) # 是否进行闭合运算以找到更多小的框
  83. more_h_lines = (
  84. kwargs.get("more_h_lines", enhance_box_line) if kwargs else enhance_box_line
  85. ) # 是否调整以找到更多的横线
  86. more_v_lines = (
  87. kwargs.get("more_v_lines", enhance_box_line) if kwargs else enhance_box_line
  88. ) # 是否调整以找到更多的横线
  89. extend_line = (
  90. kwargs.get("extend_line", enhance_box_line) if kwargs else enhance_box_line
  91. ) # 是否进行线段延长使得端点连接
  92. # 是否进行旋转修正
  93. rotated_fix = kwargs.get("rotated_fix") if kwargs else True
  94. ori_shape = img.shape
  95. pred = np.uint8(pred)
  96. hpred = copy.deepcopy(pred) # 横线
  97. vpred = copy.deepcopy(pred) # 竖线
  98. whereh = np.where(hpred == 1)
  99. wherev = np.where(vpred == 2)
  100. hpred[wherev] = 0
  101. vpred[whereh] = 0
  102. hpred = cv2.resize(hpred, (ori_shape[1], ori_shape[0]))
  103. vpred = cv2.resize(vpred, (ori_shape[1], ori_shape[0]))
  104. h, w = pred.shape
  105. hors_k = int(math.sqrt(w) * 1.2)
  106. vert_k = int(math.sqrt(h) * 1.2)
  107. hkernel = cv2.getStructuringElement(cv2.MORPH_RECT, (hors_k, 1))
  108. vkernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, vert_k))
  109. vpred = cv2.morphologyEx(
  110. vpred, cv2.MORPH_CLOSE, vkernel, iterations=1
  111. ) # 先膨胀后腐蚀的过程
  112. if morph_close:
  113. hpred = cv2.morphologyEx(hpred, cv2.MORPH_CLOSE, hkernel, iterations=1)
  114. colboxes = get_table_line(vpred, axis=1, lineW=col) # 竖线
  115. rowboxes = get_table_line(hpred, axis=0, lineW=row) # 横线
  116. rboxes_row_, rboxes_col_ = [], []
  117. if more_h_lines:
  118. rboxes_row_ = adjust_lines(rowboxes, alph=h_lines_threshold, angle=angle)
  119. if more_v_lines:
  120. rboxes_col_ = adjust_lines(colboxes, alph=v_lines_threshold, angle=angle)
  121. rowboxes += rboxes_row_
  122. colboxes += rboxes_col_
  123. if extend_line:
  124. rowboxes, colboxes = final_adjust_lines(rowboxes, colboxes)
  125. line_img = np.zeros(img.shape[:2], dtype="uint8")
  126. line_img = draw_lines(line_img, rowboxes + colboxes, color=255, lineW=2)
  127. rotated_angle = self.cal_rotate_angle(line_img)
  128. if rotated_fix and abs(rotated_angle) > 0.3:
  129. rotated_line_img = self.rotate_image(line_img, rotated_angle)
  130. rotated_polygons = self.cal_region_boxes(rotated_line_img)
  131. polygons = self.unrotate_polygons(
  132. rotated_polygons, rotated_angle, line_img.shape
  133. )
  134. else:
  135. polygons = self.cal_region_boxes(line_img)
  136. rotated_polygons = polygons.copy()
  137. return polygons, rotated_polygons
  138. def cal_region_boxes(self, tmp):
  139. labels = measure.label(tmp < 255, connectivity=2) # 8连通区域标记
  140. regions = measure.regionprops(labels)
  141. ceilboxes = min_area_rect_box(
  142. regions,
  143. False,
  144. tmp.shape[1],
  145. tmp.shape[0],
  146. filtersmall=True,
  147. adjust_box=False,
  148. ) # 最后一个参数改为False
  149. return np.array(ceilboxes)
  150. def cal_rotate_angle(self, tmp):
  151. # 计算最外侧的旋转框
  152. contours, _ = cv2.findContours(tmp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  153. if not contours:
  154. return 0
  155. largest_contour = max(contours, key=cv2.contourArea)
  156. rect = cv2.minAreaRect(largest_contour)
  157. # 计算旋转角度
  158. angle = rect[2]
  159. if angle < -45:
  160. angle += 90
  161. elif angle > 45:
  162. angle -= 90
  163. return angle
  164. def rotate_image(self, image, angle):
  165. # 获取图像的中心点
  166. (h, w) = image.shape[:2]
  167. center = (w // 2, h // 2)
  168. # 计算旋转矩阵
  169. M = cv2.getRotationMatrix2D(center, angle, 1.0)
  170. # 进行旋转
  171. rotated_image = cv2.warpAffine(
  172. image, M, (w, h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE
  173. )
  174. return rotated_image
  175. def unrotate_polygons(
  176. self, polygons: np.ndarray, angle: float, img_shape: tuple
  177. ) -> np.ndarray:
  178. # 将多边形旋转回原始位置
  179. (h, w) = img_shape
  180. center = (w // 2, h // 2)
  181. M_inv = cv2.getRotationMatrix2D(center, -angle, 1.0)
  182. # 将 (N, 8) 转换为 (N, 4, 2)
  183. polygons_reshaped = polygons.reshape(-1, 4, 2)
  184. # 批量逆旋转
  185. unrotated_polygons = cv2.transform(polygons_reshaped, M_inv)
  186. # 将 (N, 4, 2) 转换回 (N, 8)
  187. unrotated_polygons = unrotated_polygons.reshape(-1, 8)
  188. return unrotated_polygons