main.py 6.1 KB

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