main.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235
  1. # -*- encoding: utf-8 -*-
  2. # @Author: SWHL
  3. # @Contact: liekkaskono@163.com
  4. import argparse
  5. import copy
  6. import importlib
  7. import time
  8. from dataclasses import asdict, dataclass
  9. from enum import Enum
  10. from pathlib import Path
  11. from typing import Dict, List, Optional, Tuple, Union
  12. import cv2
  13. import numpy as np
  14. from rapid_table.utils import DownloadModel, LoadImage, Logger, VisTable
  15. from .matcher import TableMatch
  16. from .table_structure import TableStructurer
  17. from .table_structure_unitable import TableStructureUnitable
  18. logger = Logger(logger_name=__name__).get_log()
  19. root_dir = Path(__file__).resolve().parent
  20. class ModelType(Enum):
  21. PPSTRUCTURE_EN = "ppstructure_en"
  22. PPSTRUCTURE_ZH = "ppstructure_zh"
  23. SLANETPLUS = "slanet_plus"
  24. UNITABLE = "unitable"
  25. ROOT_URL = "https://www.modelscope.cn/models/RapidAI/RapidTable/resolve/master/"
  26. KEY_TO_MODEL_URL = {
  27. ModelType.PPSTRUCTURE_EN.value: f"{ROOT_URL}/en_ppstructure_mobile_v2_SLANet.onnx",
  28. ModelType.PPSTRUCTURE_ZH.value: f"{ROOT_URL}/ch_ppstructure_mobile_v2_SLANet.onnx",
  29. ModelType.SLANETPLUS.value: f"{ROOT_URL}/slanet-plus.onnx",
  30. ModelType.UNITABLE.value: {
  31. "encoder": f"{ROOT_URL}/unitable/encoder.pth",
  32. "decoder": f"{ROOT_URL}/unitable/decoder.pth",
  33. "vocab": f"{ROOT_URL}/unitable/vocab.json",
  34. },
  35. }
  36. @dataclass
  37. class RapidTableInput:
  38. model_type: Optional[str] = ModelType.SLANETPLUS.value
  39. model_path: Union[str, Path, None, Dict[str, str]] = None
  40. use_cuda: bool = False
  41. device: str = "cpu"
  42. @dataclass
  43. class RapidTableOutput:
  44. pred_html: Optional[str] = None
  45. cell_bboxes: Optional[np.ndarray] = None
  46. logic_points: Optional[np.ndarray] = None
  47. elapse: Optional[float] = None
  48. class RapidTable:
  49. def __init__(self, config: RapidTableInput):
  50. self.model_type = config.model_type
  51. if self.model_type not in KEY_TO_MODEL_URL:
  52. model_list = ",".join(KEY_TO_MODEL_URL)
  53. raise ValueError(
  54. f"{self.model_type} is not supported. The currently supported models are {model_list}."
  55. )
  56. config.model_path = self.get_model_path(config.model_type, config.model_path)
  57. if self.model_type == ModelType.UNITABLE.value:
  58. self.table_structure = TableStructureUnitable(asdict(config))
  59. else:
  60. self.table_structure = TableStructurer(asdict(config))
  61. self.table_matcher = TableMatch()
  62. try:
  63. self.ocr_engine = importlib.import_module("rapidocr").RapidOCR()
  64. except ModuleNotFoundError:
  65. self.ocr_engine = None
  66. self.load_img = LoadImage()
  67. def __call__(
  68. self,
  69. img_content: Union[str, np.ndarray, bytes, Path],
  70. ocr_result: List[Union[List[List[float]], str, str]] = None,
  71. ) -> RapidTableOutput:
  72. if self.ocr_engine is None and ocr_result is None:
  73. raise ValueError(
  74. "One of two conditions must be met: ocr_result is not empty, or rapidocr is installed."
  75. )
  76. img = self.load_img(img_content)
  77. s = time.perf_counter()
  78. h, w = img.shape[:2]
  79. if ocr_result is None:
  80. ocr_result = self.ocr_engine(img)
  81. ocr_result = list(
  82. zip(
  83. ocr_result.boxes,
  84. ocr_result.txts,
  85. ocr_result.scores,
  86. )
  87. )
  88. dt_boxes, rec_res = self.get_boxes_recs(ocr_result, h, w)
  89. pred_structures, cell_bboxes, _ = self.table_structure(copy.deepcopy(img))
  90. # 适配slanet-plus模型输出的box缩放还原
  91. if self.model_type == ModelType.SLANETPLUS.value:
  92. cell_bboxes = self.adapt_slanet_plus(img, cell_bboxes)
  93. pred_html = self.table_matcher(pred_structures, cell_bboxes, dt_boxes, rec_res)
  94. # 过滤掉占位的bbox
  95. mask = ~np.all(cell_bboxes == 0, axis=1)
  96. cell_bboxes = cell_bboxes[mask]
  97. logic_points = self.table_matcher.decode_logic_points(pred_structures)
  98. elapse = time.perf_counter() - s
  99. return RapidTableOutput(pred_html, cell_bboxes, logic_points, elapse)
  100. def get_boxes_recs(
  101. self, ocr_result: List[Union[List[List[float]], str, str]], h: int, w: int
  102. ) -> Tuple[np.ndarray, Tuple[str, str]]:
  103. dt_boxes, rec_res, scores = list(zip(*ocr_result))
  104. rec_res = list(zip(rec_res, scores))
  105. r_boxes = []
  106. for box in dt_boxes:
  107. box = np.array(box)
  108. x_min = max(0, box[:, 0].min() - 1)
  109. x_max = min(w, box[:, 0].max() + 1)
  110. y_min = max(0, box[:, 1].min() - 1)
  111. y_max = min(h, box[:, 1].max() + 1)
  112. box = [x_min, y_min, x_max, y_max]
  113. r_boxes.append(box)
  114. dt_boxes = np.array(r_boxes)
  115. return dt_boxes, rec_res
  116. def adapt_slanet_plus(self, img: np.ndarray, cell_bboxes: np.ndarray) -> np.ndarray:
  117. h, w = img.shape[:2]
  118. resized = 488
  119. ratio = min(resized / h, resized / w)
  120. w_ratio = resized / (w * ratio)
  121. h_ratio = resized / (h * ratio)
  122. cell_bboxes[:, 0::2] *= w_ratio
  123. cell_bboxes[:, 1::2] *= h_ratio
  124. return cell_bboxes
  125. @staticmethod
  126. def get_model_path(
  127. model_type: str, model_path: Union[str, Path, None]
  128. ) -> Union[str, Dict[str, str]]:
  129. if model_path is not None:
  130. return model_path
  131. model_url = KEY_TO_MODEL_URL.get(model_type, None)
  132. if isinstance(model_url, str):
  133. model_path = DownloadModel.download(model_url)
  134. return model_path
  135. if isinstance(model_url, dict):
  136. model_paths = {}
  137. for k, url in model_url.items():
  138. model_paths[k] = DownloadModel.download(
  139. url, save_model_name=f"{model_type}_{Path(url).name}"
  140. )
  141. return model_paths
  142. raise ValueError(f"Model URL: {type(model_url)} is not between str and dict.")
  143. def parse_args(arg_list: Optional[List[str]] = None):
  144. parser = argparse.ArgumentParser()
  145. parser.add_argument(
  146. "-v",
  147. "--vis",
  148. action="store_true",
  149. default=False,
  150. help="Wheter to visualize the layout results.",
  151. )
  152. parser.add_argument(
  153. "-img", "--img_path", type=str, required=True, help="Path to image for layout."
  154. )
  155. parser.add_argument(
  156. "-m",
  157. "--model_type",
  158. type=str,
  159. default=ModelType.SLANETPLUS.value,
  160. choices=list(KEY_TO_MODEL_URL),
  161. )
  162. args = parser.parse_args(arg_list)
  163. return args
  164. def main(arg_list: Optional[List[str]] = None):
  165. args = parse_args(arg_list)
  166. try:
  167. ocr_engine = importlib.import_module("rapidocr").RapidOCR()
  168. except ModuleNotFoundError as exc:
  169. raise ModuleNotFoundError(
  170. "Please install the rapidocr by pip install rapidocr"
  171. ) from exc
  172. input_args = RapidTableInput(model_type=args.model_type)
  173. table_engine = RapidTable(input_args)
  174. img = cv2.imread(args.img_path)
  175. rapid_ocr_output = ocr_engine(img)
  176. ocr_result = list(
  177. zip(rapid_ocr_output.boxes, rapid_ocr_output.txts, rapid_ocr_output.scores)
  178. )
  179. table_results = table_engine(img, ocr_result)
  180. print(table_results.pred_html)
  181. viser = VisTable()
  182. if args.vis:
  183. img_path = Path(args.img_path)
  184. save_dir = img_path.resolve().parent
  185. save_html_path = save_dir / f"{Path(img_path).stem}.html"
  186. save_drawed_path = save_dir / f"vis_{Path(img_path).name}"
  187. viser(img_path, table_results, save_html_path, save_drawed_path)
  188. if __name__ == "__main__":
  189. main()