wired_table_rec_utils.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565
  1. # -*- encoding: utf-8 -*-
  2. import math
  3. import os
  4. import platform
  5. import traceback
  6. from enum import Enum
  7. from io import BytesIO
  8. from pathlib import Path
  9. from typing import List, Union, Dict, Any, Tuple, Optional
  10. import cv2
  11. import numpy as np
  12. from onnxruntime import (
  13. GraphOptimizationLevel,
  14. InferenceSession,
  15. SessionOptions,
  16. get_available_providers,
  17. get_device,
  18. )
  19. from PIL import Image, UnidentifiedImageError
  20. root_dir = Path(__file__).resolve().parent
  21. InputType = Union[str, np.ndarray, bytes, Path]
  22. class EP(Enum):
  23. CPU_EP = "CPUExecutionProvider"
  24. class OrtInferSession:
  25. def __init__(self, config: Dict[str, Any]):
  26. model_path = config.get("model_path", None)
  27. self._verify_model(model_path)
  28. self.had_providers: List[str] = get_available_providers()
  29. EP_list = self._get_ep_list()
  30. sess_opt = self._init_sess_opts(config)
  31. self.session = InferenceSession(
  32. model_path,
  33. sess_options=sess_opt,
  34. providers=EP_list,
  35. )
  36. @staticmethod
  37. def _init_sess_opts(config: Dict[str, Any]) -> SessionOptions:
  38. sess_opt = SessionOptions()
  39. sess_opt.log_severity_level = 4
  40. sess_opt.enable_cpu_mem_arena = False
  41. sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
  42. cpu_nums = os.cpu_count()
  43. intra_op_num_threads = config.get("intra_op_num_threads", -1)
  44. if intra_op_num_threads != -1 and 1 <= intra_op_num_threads <= cpu_nums:
  45. sess_opt.intra_op_num_threads = intra_op_num_threads
  46. inter_op_num_threads = config.get("inter_op_num_threads", -1)
  47. if inter_op_num_threads != -1 and 1 <= inter_op_num_threads <= cpu_nums:
  48. sess_opt.inter_op_num_threads = inter_op_num_threads
  49. return sess_opt
  50. def get_metadata(self, key: str = "character") -> list:
  51. meta_dict = self.session.get_modelmeta().custom_metadata_map
  52. content_list = meta_dict[key].splitlines()
  53. return content_list
  54. def _get_ep_list(self) -> List[Tuple[str, Dict[str, Any]]]:
  55. cpu_provider_opts = {
  56. "arena_extend_strategy": "kSameAsRequested",
  57. }
  58. EP_list = [(EP.CPU_EP.value, cpu_provider_opts)]
  59. return EP_list
  60. def __call__(self, input_content: List[np.ndarray]) -> np.ndarray:
  61. input_dict = dict(zip(self.get_input_names(), input_content))
  62. try:
  63. return self.session.run(None, input_dict)
  64. except Exception as e:
  65. error_info = traceback.format_exc()
  66. raise ONNXRuntimeError(error_info) from e
  67. def get_input_names(self) -> List[str]:
  68. return [v.name for v in self.session.get_inputs()]
  69. def get_output_names(self) -> List[str]:
  70. return [v.name for v in self.session.get_outputs()]
  71. def get_character_list(self, key: str = "character") -> List[str]:
  72. meta_dict = self.session.get_modelmeta().custom_metadata_map
  73. return meta_dict[key].splitlines()
  74. def have_key(self, key: str = "character") -> bool:
  75. meta_dict = self.session.get_modelmeta().custom_metadata_map
  76. if key in meta_dict.keys():
  77. return True
  78. return False
  79. @staticmethod
  80. def _verify_model(model_path: Union[str, Path, None]):
  81. if model_path is None:
  82. raise ValueError("model_path is None!")
  83. model_path = Path(model_path)
  84. if not model_path.exists():
  85. raise FileNotFoundError(f"{model_path} does not exists.")
  86. if not model_path.is_file():
  87. raise FileExistsError(f"{model_path} is not a file.")
  88. class ONNXRuntimeError(Exception):
  89. pass
  90. class LoadImage:
  91. def __init__(
  92. self,
  93. ):
  94. pass
  95. def __call__(self, img: InputType) -> np.ndarray:
  96. if not isinstance(img, InputType.__args__):
  97. raise LoadImageError(
  98. f"The img type {type(img)} does not in {InputType.__args__}"
  99. )
  100. img = self.load_img(img)
  101. img = self.convert_img(img)
  102. return img
  103. def load_img(self, img: InputType) -> np.ndarray:
  104. if isinstance(img, (str, Path)):
  105. self.verify_exist(img)
  106. try:
  107. img = np.array(Image.open(img))
  108. except UnidentifiedImageError as e:
  109. raise LoadImageError(f"cannot identify image file {img}") from e
  110. return img
  111. if isinstance(img, bytes):
  112. img = np.array(Image.open(BytesIO(img)))
  113. return img
  114. if isinstance(img, np.ndarray):
  115. return img
  116. raise LoadImageError(f"{type(img)} is not supported!")
  117. def convert_img(self, img: np.ndarray):
  118. if img.ndim == 2:
  119. return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
  120. if img.ndim == 3:
  121. channel = img.shape[2]
  122. if channel == 1:
  123. return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
  124. if channel == 2:
  125. return self.cvt_two_to_three(img)
  126. if channel == 4:
  127. return self.cvt_four_to_three(img)
  128. if channel == 3:
  129. return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
  130. raise LoadImageError(
  131. f"The channel({channel}) of the img is not in [1, 2, 3, 4]"
  132. )
  133. raise LoadImageError(f"The ndim({img.ndim}) of the img is not in [2, 3]")
  134. @staticmethod
  135. def cvt_four_to_three(img: np.ndarray) -> np.ndarray:
  136. """RGBA → BGR"""
  137. r, g, b, a = cv2.split(img)
  138. new_img = cv2.merge((b, g, r))
  139. not_a = cv2.bitwise_not(a)
  140. not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
  141. new_img = cv2.bitwise_and(new_img, new_img, mask=a)
  142. new_img = cv2.add(new_img, not_a)
  143. return new_img
  144. @staticmethod
  145. def cvt_two_to_three(img: np.ndarray) -> np.ndarray:
  146. """gray + alpha → BGR"""
  147. img_gray = img[..., 0]
  148. img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR)
  149. img_alpha = img[..., 1]
  150. not_a = cv2.bitwise_not(img_alpha)
  151. not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
  152. new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha)
  153. new_img = cv2.add(new_img, not_a)
  154. return new_img
  155. @staticmethod
  156. def verify_exist(file_path: Union[str, Path]):
  157. if not Path(file_path).exists():
  158. raise LoadImageError(f"{file_path} does not exist.")
  159. class LoadImageError(Exception):
  160. pass
  161. # Pillow >=v9.1.0 use a slightly different naming scheme for filters.
  162. # Set pillow_interp_codes according to the naming scheme used.
  163. if Image is not None:
  164. if hasattr(Image, "Resampling"):
  165. pillow_interp_codes = {
  166. "nearest": Image.Resampling.NEAREST,
  167. "bilinear": Image.Resampling.BILINEAR,
  168. "bicubic": Image.Resampling.BICUBIC,
  169. "box": Image.Resampling.BOX,
  170. "lanczos": Image.Resampling.LANCZOS,
  171. "hamming": Image.Resampling.HAMMING,
  172. }
  173. else:
  174. pillow_interp_codes = {
  175. "nearest": Image.NEAREST,
  176. "bilinear": Image.BILINEAR,
  177. "bicubic": Image.BICUBIC,
  178. "box": Image.BOX,
  179. "lanczos": Image.LANCZOS,
  180. "hamming": Image.HAMMING,
  181. }
  182. cv2_interp_codes = {
  183. "nearest": cv2.INTER_NEAREST,
  184. "bilinear": cv2.INTER_LINEAR,
  185. "bicubic": cv2.INTER_CUBIC,
  186. "area": cv2.INTER_AREA,
  187. "lanczos": cv2.INTER_LANCZOS4,
  188. }
  189. def resize_img(img, scale, keep_ratio=True):
  190. if keep_ratio:
  191. # 缩小使用area更保真
  192. if min(img.shape[:2]) > min(scale):
  193. interpolation = "area"
  194. else:
  195. interpolation = "bicubic" # bilinear
  196. img_new, scale_factor = imrescale(
  197. img, scale, return_scale=True, interpolation=interpolation
  198. )
  199. # the w_scale and h_scale has minor difference
  200. # a real fix should be done in the mmcv.imrescale in the future
  201. new_h, new_w = img_new.shape[:2]
  202. h, w = img.shape[:2]
  203. w_scale = new_w / w
  204. h_scale = new_h / h
  205. else:
  206. img_new, w_scale, h_scale = imresize(img, scale, return_scale=True)
  207. return img_new, w_scale, h_scale
  208. def imrescale(img, scale, return_scale=False, interpolation="bilinear", backend=None):
  209. """Resize image while keeping the aspect ratio.
  210. Args:
  211. img (ndarray): The input image.
  212. scale (float | tuple[int]): The scaling factor or maximum size.
  213. If it is a float number, then the image will be rescaled by this
  214. factor, else if it is a tuple of 2 integers, then the image will
  215. be rescaled as large as possible within the scale.
  216. return_scale (bool): Whether to return the scaling factor besides the
  217. rescaled image.
  218. interpolation (str): Same as :func:`resize`.
  219. backend (str | None): Same as :func:`resize`.
  220. Returns:
  221. ndarray: The rescaled image.
  222. """
  223. h, w = img.shape[:2]
  224. new_size, scale_factor = rescale_size((w, h), scale, return_scale=True)
  225. rescaled_img = imresize(img, new_size, interpolation=interpolation, backend=backend)
  226. if return_scale:
  227. return rescaled_img, scale_factor
  228. else:
  229. return rescaled_img
  230. def imresize(
  231. img, size, return_scale=False, interpolation="bilinear", out=None, backend=None
  232. ):
  233. """Resize image to a given size.
  234. Args:
  235. img (ndarray): The input image.
  236. size (tuple[int]): Target size (w, h).
  237. return_scale (bool): Whether to return `w_scale` and `h_scale`.
  238. interpolation (str): Interpolation method, accepted values are
  239. "nearest", "bilinear", "bicubic", "area", "lanczos" for 'cv2'
  240. backend, "nearest", "bilinear" for 'pillow' backend.
  241. out (ndarray): The output destination.
  242. backend (str | None): The image resize backend type. Options are `cv2`,
  243. `pillow`, `None`. If backend is None, the global imread_backend
  244. specified by ``mmcv.use_backend()`` will be used. Default: None.
  245. Returns:
  246. tuple | ndarray: (`resized_img`, `w_scale`, `h_scale`) or
  247. `resized_img`.
  248. """
  249. h, w = img.shape[:2]
  250. if backend is None:
  251. backend = "cv2"
  252. if backend not in ["cv2", "pillow"]:
  253. raise ValueError(
  254. f"backend: {backend} is not supported for resize."
  255. f"Supported backends are 'cv2', 'pillow'"
  256. )
  257. if backend == "pillow":
  258. assert img.dtype == np.uint8, "Pillow backend only support uint8 type"
  259. pil_image = Image.fromarray(img)
  260. pil_image = pil_image.resize(size, pillow_interp_codes[interpolation])
  261. resized_img = np.array(pil_image)
  262. else:
  263. resized_img = cv2.resize(
  264. img, size, dst=out, interpolation=cv2_interp_codes[interpolation]
  265. )
  266. if not return_scale:
  267. return resized_img
  268. else:
  269. w_scale = size[0] / w
  270. h_scale = size[1] / h
  271. return resized_img, w_scale, h_scale
  272. def rescale_size(old_size, scale, return_scale=False):
  273. """Calculate the new size to be rescaled to.
  274. Args:
  275. old_size (tuple[int]): The old size (w, h) of image.
  276. scale (float | tuple[int]): The scaling factor or maximum size.
  277. If it is a float number, then the image will be rescaled by this
  278. factor, else if it is a tuple of 2 integers, then the image will
  279. be rescaled as large as possible within the scale.
  280. return_scale (bool): Whether to return the scaling factor besides the
  281. rescaled image size.
  282. Returns:
  283. tuple[int]: The new rescaled image size.
  284. """
  285. w, h = old_size
  286. if isinstance(scale, (float, int)):
  287. if scale <= 0:
  288. raise ValueError(f"Invalid scale {scale}, must be positive.")
  289. scale_factor = scale
  290. elif isinstance(scale, tuple):
  291. max_long_edge = max(scale)
  292. max_short_edge = min(scale)
  293. scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
  294. else:
  295. raise TypeError(
  296. f"Scale must be a number or tuple of int, but got {type(scale)}"
  297. )
  298. new_size = _scale_size((w, h), scale_factor)
  299. if return_scale:
  300. return new_size, scale_factor
  301. else:
  302. return new_size
  303. def _scale_size(size, scale):
  304. """Rescale a size by a ratio.
  305. Args:
  306. size (tuple[int]): (w, h).
  307. scale (float | tuple(float)): Scaling factor.
  308. Returns:
  309. tuple[int]: scaled size.
  310. """
  311. if isinstance(scale, (float, int)):
  312. scale = (scale, scale)
  313. w, h = size
  314. return int(w * float(scale[0]) + 0.5), int(h * float(scale[1]) + 0.5)
  315. class ImageOrientationCorrector:
  316. """
  317. 对图片小角度(-90 - + 90度进行修正)
  318. """
  319. def __init__(self):
  320. self.img_loader = LoadImage()
  321. def __call__(self, img: InputType):
  322. img = self.img_loader(img)
  323. # 取灰度
  324. gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  325. # 二值化
  326. gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
  327. # 边缘检测
  328. edges = cv2.Canny(gray, 100, 250, apertureSize=3)
  329. # 霍夫变换,摘自https://blog.csdn.net/feilong_csdn/article/details/81586322
  330. lines = cv2.HoughLines(edges, 1, np.pi / 180, 0)
  331. for rho, theta in lines[0]:
  332. a = np.cos(theta)
  333. b = np.sin(theta)
  334. x0 = a * rho
  335. y0 = b * rho
  336. x1 = int(x0 + 1000 * (-b))
  337. y1 = int(y0 + 1000 * (a))
  338. x2 = int(x0 - 1000 * (-b))
  339. y2 = int(y0 - 1000 * (a))
  340. if x1 == x2 or y1 == y2:
  341. return img
  342. else:
  343. t = float(y2 - y1) / (x2 - x1)
  344. # 得到角度后
  345. rotate_angle = math.degrees(math.atan(t))
  346. if rotate_angle > 45:
  347. rotate_angle = -90 + rotate_angle
  348. elif rotate_angle < -45:
  349. rotate_angle = 90 + rotate_angle
  350. # 旋转图像
  351. (h, w) = img.shape[:2]
  352. center = (w // 2, h // 2)
  353. M = cv2.getRotationMatrix2D(center, rotate_angle, 1.0)
  354. return cv2.warpAffine(img, M, (w, h))
  355. class VisTable:
  356. def __init__(self):
  357. self.load_img = LoadImage()
  358. def __call__(
  359. self,
  360. img_path: Union[str, Path],
  361. table_results,
  362. save_html_path: Optional[Union[str, Path]] = None,
  363. save_drawed_path: Optional[Union[str, Path]] = None,
  364. save_logic_path: Optional[Union[str, Path]] = None,
  365. ):
  366. if save_html_path:
  367. html_with_border = self.insert_border_style(table_results.pred_html)
  368. self.save_html(save_html_path, html_with_border)
  369. table_cell_bboxes = table_results.cell_bboxes
  370. table_logic_points = table_results.logic_points
  371. if table_cell_bboxes is None:
  372. return None
  373. img = self.load_img(img_path)
  374. dims_bboxes = table_cell_bboxes.shape[1]
  375. if dims_bboxes == 4:
  376. drawed_img = self.draw_rectangle(img, table_cell_bboxes)
  377. elif dims_bboxes == 8:
  378. drawed_img = self.draw_polylines(img, table_cell_bboxes)
  379. else:
  380. raise ValueError("Shape of table bounding boxes is not between in 4 or 8.")
  381. if save_drawed_path:
  382. self.save_img(save_drawed_path, drawed_img)
  383. if save_logic_path:
  384. polygons = [[box[0], box[1], box[4], box[5]] for box in table_cell_bboxes]
  385. self.plot_rec_box_with_logic_info(
  386. img_path, save_logic_path, table_logic_points, polygons
  387. )
  388. return drawed_img
  389. def insert_border_style(self, table_html_str: str):
  390. style_res = """<meta charset="UTF-8"><style>
  391. table {
  392. border-collapse: collapse;
  393. width: 100%;
  394. }
  395. th, td {
  396. border: 1px solid black;
  397. padding: 8px;
  398. text-align: center;
  399. }
  400. th {
  401. background-color: #f2f2f2;
  402. }
  403. </style>"""
  404. prefix_table, suffix_table = table_html_str.split("<body>")
  405. html_with_border = f"{prefix_table}{style_res}<body>{suffix_table}"
  406. return html_with_border
  407. def plot_rec_box_with_logic_info(
  408. self, img_path, output_path, logic_points, sorted_polygons
  409. ):
  410. """
  411. :param img_path
  412. :param output_path
  413. :param logic_points: [row_start,row_end,col_start,col_end]
  414. :param sorted_polygons: [xmin,ymin,xmax,ymax]
  415. :return:
  416. """
  417. # 读取原图
  418. img = cv2.imread(img_path)
  419. img = cv2.copyMakeBorder(
  420. img, 0, 0, 0, 100, cv2.BORDER_CONSTANT, value=[255, 255, 255]
  421. )
  422. # 绘制 polygons 矩形
  423. for idx, polygon in enumerate(sorted_polygons):
  424. x0, y0, x1, y1 = polygon[0], polygon[1], polygon[2], polygon[3]
  425. x0 = round(x0)
  426. y0 = round(y0)
  427. x1 = round(x1)
  428. y1 = round(y1)
  429. cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), 1)
  430. # 增大字体大小和线宽
  431. font_scale = 0.9 # 原先是0.5
  432. thickness = 1 # 原先是1
  433. logic_point = logic_points[idx]
  434. cv2.putText(
  435. img,
  436. f"row: {logic_point[0]}-{logic_point[1]}",
  437. (x0 + 3, y0 + 8),
  438. cv2.FONT_HERSHEY_PLAIN,
  439. font_scale,
  440. (0, 0, 255),
  441. thickness,
  442. )
  443. cv2.putText(
  444. img,
  445. f"col: {logic_point[2]}-{logic_point[3]}",
  446. (x0 + 3, y0 + 18),
  447. cv2.FONT_HERSHEY_PLAIN,
  448. font_scale,
  449. (0, 0, 255),
  450. thickness,
  451. )
  452. os.makedirs(os.path.dirname(output_path), exist_ok=True)
  453. # 保存绘制后的图像
  454. self.save_img(output_path, img)
  455. @staticmethod
  456. def draw_rectangle(img: np.ndarray, boxes: np.ndarray) -> np.ndarray:
  457. img_copy = img.copy()
  458. for box in boxes.astype(int):
  459. x1, y1, x2, y2 = box
  460. cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)
  461. return img_copy
  462. @staticmethod
  463. def draw_polylines(img: np.ndarray, points) -> np.ndarray:
  464. img_copy = img.copy()
  465. for point in points.astype(int):
  466. point = point.reshape(4, 2)
  467. cv2.polylines(img_copy, [point.astype(int)], True, (255, 0, 0), 2)
  468. return img_copy
  469. @staticmethod
  470. def save_img(save_path: Union[str, Path], img: np.ndarray):
  471. cv2.imwrite(str(save_path), img)
  472. @staticmethod
  473. def save_html(save_path: Union[str, Path], html: str):
  474. with open(save_path, "w", encoding="utf-8") as f:
  475. f.write(html)