pipeline_v2.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os, sys
  15. from typing import Any, Dict, Optional, Union, List, Tuple
  16. import numpy as np
  17. import cv2
  18. from ..base import BasePipeline
  19. from ..components import CropByBoxes
  20. from .utils import get_neighbor_boxes_idx
  21. from .table_recognition_post_processing_v2 import get_table_recognition_res
  22. from .result import SingleTableRecognitionResult, TableRecognitionResult
  23. from ....utils import logging
  24. from ...utils.pp_option import PaddlePredictorOption
  25. from ...common.reader import ReadImage
  26. from ...common.batch_sampler import ImageBatchSampler
  27. from ..ocr.result import OCRResult
  28. from ..doc_preprocessor.result import DocPreprocessorResult
  29. from ...models.object_detection.result import DetResult
  30. class TableRecognitionPipelineV2(BasePipeline):
  31. """Table Recognition Pipeline"""
  32. entities = ["table_recognition_v2"]
  33. def __init__(
  34. self,
  35. config: Dict,
  36. device: str = None,
  37. pp_option: PaddlePredictorOption = None,
  38. use_hpip: bool = False,
  39. hpi_params: Optional[Dict[str, Any]] = None,
  40. ) -> None:
  41. """Initializes the layout parsing pipeline.
  42. Args:
  43. config (Dict): Configuration dictionary containing various settings.
  44. device (str, optional): Device to run the predictions on. Defaults to None.
  45. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  46. use_hpip (bool, optional): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
  47. hpi_params (Optional[Dict[str, Any]], optional): HPIP parameters. Defaults to None.
  48. """
  49. super().__init__(
  50. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_params=hpi_params
  51. )
  52. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  53. if self.use_doc_preprocessor:
  54. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  55. "DocPreprocessor",
  56. {
  57. "pipeline_config_error": "config error for doc_preprocessor_pipeline!"
  58. },
  59. )
  60. self.doc_preprocessor_pipeline = self.create_pipeline(
  61. doc_preprocessor_config
  62. )
  63. self.use_layout_detection = config.get("use_layout_detection", True)
  64. if self.use_layout_detection:
  65. layout_det_config = config.get("SubModules", {}).get(
  66. "LayoutDetection",
  67. {"model_config_error": "config error for layout_det_model!"},
  68. )
  69. self.layout_det_model = self.create_model(layout_det_config)
  70. table_cls_config = config.get("SubModules", {}).get(
  71. "TableClassification",
  72. {"model_config_error": "config error for table_classification_model!"},
  73. )
  74. self.table_cls_model = self.create_model(table_cls_config)
  75. wired_table_rec_config = config.get("SubModules", {}).get(
  76. "WiredTableStructureRecognition",
  77. {"model_config_error": "config error for wired_table_structure_model!"},
  78. )
  79. self.wired_table_rec_model = self.create_model(wired_table_rec_config)
  80. wireless_table_rec_config = config.get("SubModules", {}).get(
  81. "WirelessTableStructureRecognition",
  82. {"model_config_error": "config error for wireless_table_structure_model!"},
  83. )
  84. self.wireless_table_rec_model = self.create_model(wireless_table_rec_config)
  85. wired_table_cells_det_config = config.get("SubModules", {}).get(
  86. "WiredTableCellsDetection",
  87. {
  88. "model_config_error": "config error for wired_table_cells_detection_model!"
  89. },
  90. )
  91. self.wired_table_cells_detection_model = self.create_model(
  92. wired_table_cells_det_config
  93. )
  94. wireless_table_cells_det_config = config.get("SubModules", {}).get(
  95. "WirelessTableCellsDetection",
  96. {
  97. "model_config_error": "config error for wireless_table_cells_detection_model!"
  98. },
  99. )
  100. self.wireless_table_cells_detection_model = self.create_model(
  101. wireless_table_cells_det_config
  102. )
  103. self.use_ocr_model = config.get("use_ocr_model", True)
  104. if self.use_ocr_model:
  105. general_ocr_config = config.get("SubPipelines", {}).get(
  106. "GeneralOCR",
  107. {"pipeline_config_error": "config error for general_ocr_pipeline!"},
  108. )
  109. self.general_ocr_pipeline = self.create_pipeline(general_ocr_config)
  110. self._crop_by_boxes = CropByBoxes()
  111. self.batch_sampler = ImageBatchSampler(batch_size=1)
  112. self.img_reader = ReadImage(format="BGR")
  113. def get_model_settings(
  114. self,
  115. use_doc_orientation_classify: Optional[bool],
  116. use_doc_unwarping: Optional[bool],
  117. use_layout_detection: Optional[bool],
  118. use_ocr_model: Optional[bool],
  119. ) -> dict:
  120. """
  121. Get the model settings based on the provided parameters or default values.
  122. Args:
  123. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  124. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  125. use_layout_detection (Optional[bool]): Whether to use layout detection.
  126. use_ocr_model (Optional[bool]): Whether to use OCR model.
  127. Returns:
  128. dict: A dictionary containing the model settings.
  129. """
  130. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  131. use_doc_preprocessor = self.use_doc_preprocessor
  132. else:
  133. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  134. use_doc_preprocessor = True
  135. else:
  136. use_doc_preprocessor = False
  137. if use_layout_detection is None:
  138. use_layout_detection = self.use_layout_detection
  139. if use_ocr_model is None:
  140. use_ocr_model = self.use_ocr_model
  141. return dict(
  142. use_doc_preprocessor=use_doc_preprocessor,
  143. use_layout_detection=use_layout_detection,
  144. use_ocr_model=use_ocr_model,
  145. )
  146. def check_model_settings_valid(
  147. self,
  148. model_settings: Dict,
  149. overall_ocr_res: OCRResult,
  150. layout_det_res: DetResult,
  151. ) -> bool:
  152. """
  153. Check if the input parameters are valid based on the initialized models.
  154. Args:
  155. model_settings (Dict): A dictionary containing input parameters.
  156. overall_ocr_res (OCRResult): Overall OCR result obtained after running the OCR pipeline.
  157. The overall OCR result with convert_points_to_boxes information.
  158. layout_det_res (DetResult): The layout detection result.
  159. Returns:
  160. bool: True if all required models are initialized according to input parameters, False otherwise.
  161. """
  162. if model_settings["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  163. logging.error(
  164. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  165. )
  166. return False
  167. if model_settings["use_layout_detection"]:
  168. if layout_det_res is not None:
  169. logging.error(
  170. "The layout detection model has already been initialized, please set use_layout_detection=False"
  171. )
  172. return False
  173. if not self.use_layout_detection:
  174. logging.error(
  175. "Set use_layout_detection, but the models for layout detection are not initialized."
  176. )
  177. return False
  178. if model_settings["use_ocr_model"]:
  179. if overall_ocr_res is not None:
  180. logging.error(
  181. "The OCR models have already been initialized, please set use_ocr_model=False"
  182. )
  183. return False
  184. if not self.use_ocr_model:
  185. logging.error(
  186. "Set use_ocr_model, but the models for OCR are not initialized."
  187. )
  188. return False
  189. else:
  190. if overall_ocr_res is None:
  191. logging.error("Set use_ocr_model=False, but no OCR results were found.")
  192. return False
  193. return True
  194. def predict_doc_preprocessor_res(
  195. self, image_array: np.ndarray, input_params: dict
  196. ) -> Tuple[DocPreprocessorResult, np.ndarray]:
  197. """
  198. Preprocess the document image based on input parameters.
  199. Args:
  200. image_array (np.ndarray): The input image array.
  201. input_params (dict): Dictionary containing preprocessing parameters.
  202. Returns:
  203. tuple[DocPreprocessorResult, np.ndarray]: A tuple containing the preprocessing
  204. result dictionary and the processed image array.
  205. """
  206. if input_params["use_doc_preprocessor"]:
  207. use_doc_orientation_classify = input_params["use_doc_orientation_classify"]
  208. use_doc_unwarping = input_params["use_doc_unwarping"]
  209. doc_preprocessor_res = next(
  210. self.doc_preprocessor_pipeline(
  211. image_array,
  212. use_doc_orientation_classify=use_doc_orientation_classify,
  213. use_doc_unwarping=use_doc_unwarping,
  214. )
  215. )
  216. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  217. else:
  218. doc_preprocessor_res = {}
  219. doc_preprocessor_image = image_array
  220. return doc_preprocessor_res, doc_preprocessor_image
  221. def extract_results(self, pred, task):
  222. if task == "cls":
  223. return pred["label_names"][np.argmax(pred["scores"])]
  224. elif task == "det":
  225. threshold = 0.0
  226. result = []
  227. if "boxes" in pred and isinstance(pred["boxes"], list):
  228. for box in pred["boxes"]:
  229. if isinstance(box, dict) and "score" in box and "coordinate" in box:
  230. score = box["score"]
  231. coordinate = box["coordinate"]
  232. if isinstance(score, float) and score > threshold:
  233. result.append(coordinate)
  234. return result
  235. elif task == "table_stru":
  236. return pred["structure"]
  237. else:
  238. return None
  239. def predict_single_table_recognition_res(
  240. self,
  241. image_array: np.ndarray,
  242. overall_ocr_res: OCRResult,
  243. table_box: list,
  244. flag_find_nei_text: bool = True,
  245. ) -> SingleTableRecognitionResult:
  246. """
  247. Predict table recognition results from an image array, layout detection results, and OCR results.
  248. Args:
  249. image_array (np.ndarray): The input image represented as a numpy array.
  250. overall_ocr_res (OCRResult): Overall OCR result obtained after running the OCR pipeline.
  251. The overall OCR results containing text recognition information.
  252. table_box (list): The table box coordinates.
  253. flag_find_nei_text (bool): Whether to find neighboring text.
  254. Returns:
  255. SingleTableRecognitionResult: single table recognition result.
  256. """
  257. table_cls_pred = next(self.table_cls_model(image_array))
  258. table_cls_result = self.extract_results(table_cls_pred, "cls")
  259. if table_cls_result == "wired_table":
  260. table_structure_pred = next(self.wired_table_rec_model(image_array))
  261. table_cells_pred = next(self.wired_table_cells_detection_model(image_array))
  262. elif table_cls_result == "wireless_table":
  263. table_structure_pred = next(self.wireless_table_rec_model(image_array))
  264. table_cells_pred = next(
  265. self.wireless_table_cells_detection_model(image_array)
  266. )
  267. table_structure_result = self.extract_results(
  268. table_structure_pred, "table_stru"
  269. )
  270. table_cells_result = self.extract_results(table_cells_pred, "det")
  271. single_table_recognition_res = get_table_recognition_res(
  272. table_box, table_structure_result, table_cells_result, overall_ocr_res
  273. )
  274. neighbor_text = ""
  275. if flag_find_nei_text:
  276. match_idx_list = get_neighbor_boxes_idx(
  277. overall_ocr_res["rec_boxes"], table_box
  278. )
  279. if len(match_idx_list) > 0:
  280. for idx in match_idx_list:
  281. neighbor_text += overall_ocr_res["rec_texts"][idx] + "; "
  282. single_table_recognition_res["neighbor_text"] = neighbor_text
  283. return single_table_recognition_res
  284. def predict(
  285. self,
  286. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  287. use_doc_orientation_classify: Optional[bool] = None,
  288. use_doc_unwarping: Optional[bool] = None,
  289. use_layout_detection: Optional[bool] = None,
  290. use_ocr_model: Optional[bool] = None,
  291. overall_ocr_res: Optional[OCRResult] = None,
  292. layout_det_res: Optional[DetResult] = None,
  293. text_det_limit_side_len: Optional[int] = None,
  294. text_det_limit_type: Optional[str] = None,
  295. text_det_thresh: Optional[float] = None,
  296. text_det_box_thresh: Optional[float] = None,
  297. text_det_unclip_ratio: Optional[float] = None,
  298. text_rec_score_thresh: Optional[float] = None,
  299. **kwargs,
  300. ) -> TableRecognitionResult:
  301. """
  302. This function predicts the layout parsing result for the given input.
  303. Args:
  304. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) of pdf(s) to be processed.
  305. use_layout_detection (bool): Whether to use layout detection.
  306. use_doc_orientation_classify (bool): Whether to use document orientation classification.
  307. use_doc_unwarping (bool): Whether to use document unwarping.
  308. overall_ocr_res (OCRResult): The overall OCR result with convert_points_to_boxes information.
  309. It will be used if it is not None and use_ocr_model is False.
  310. layout_det_res (DetResult): The layout detection result.
  311. It will be used if it is not None and use_layout_detection is False.
  312. **kwargs: Additional keyword arguments.
  313. Returns:
  314. TableRecognitionResult: The predicted table recognition result.
  315. """
  316. model_settings = self.get_model_settings(
  317. use_doc_orientation_classify,
  318. use_doc_unwarping,
  319. use_layout_detection,
  320. use_ocr_model,
  321. )
  322. if not self.check_model_settings_valid(
  323. model_settings, overall_ocr_res, layout_det_res
  324. ):
  325. yield {"error": "the input params for model settings are invalid!"}
  326. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  327. image_array = self.img_reader(batch_data.instances)[0]
  328. if model_settings["use_doc_preprocessor"]:
  329. doc_preprocessor_res = next(
  330. self.doc_preprocessor_pipeline(
  331. image_array,
  332. use_doc_orientation_classify=use_doc_orientation_classify,
  333. use_doc_unwarping=use_doc_unwarping,
  334. )
  335. )
  336. else:
  337. doc_preprocessor_res = {"output_img": image_array}
  338. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  339. if model_settings["use_ocr_model"]:
  340. overall_ocr_res = next(
  341. self.general_ocr_pipeline(
  342. doc_preprocessor_image,
  343. text_det_limit_side_len=text_det_limit_side_len,
  344. text_det_limit_type=text_det_limit_type,
  345. text_det_thresh=text_det_thresh,
  346. text_det_box_thresh=text_det_box_thresh,
  347. text_det_unclip_ratio=text_det_unclip_ratio,
  348. text_rec_score_thresh=text_rec_score_thresh,
  349. )
  350. )
  351. table_res_list = []
  352. table_region_id = 1
  353. if not model_settings["use_layout_detection"] and layout_det_res is None:
  354. layout_det_res = {}
  355. img_height, img_width = doc_preprocessor_image.shape[:2]
  356. table_box = [0, 0, img_width - 1, img_height - 1]
  357. single_table_rec_res = self.predict_single_table_recognition_res(
  358. doc_preprocessor_image,
  359. overall_ocr_res,
  360. table_box,
  361. flag_find_nei_text=False,
  362. )
  363. single_table_rec_res["table_region_id"] = table_region_id
  364. table_res_list.append(single_table_rec_res)
  365. table_region_id += 1
  366. else:
  367. if model_settings["use_layout_detection"]:
  368. layout_det_res = next(self.layout_det_model(doc_preprocessor_image))
  369. for box_info in layout_det_res["boxes"]:
  370. if box_info["label"].lower() in ["table"]:
  371. crop_img_info = self._crop_by_boxes(image_array, [box_info])
  372. crop_img_info = crop_img_info[0]
  373. table_box = crop_img_info["box"]
  374. single_table_rec_res = (
  375. self.predict_single_table_recognition_res(
  376. crop_img_info["img"], overall_ocr_res, table_box
  377. )
  378. )
  379. single_table_rec_res["table_region_id"] = table_region_id
  380. table_res_list.append(single_table_rec_res)
  381. table_region_id += 1
  382. single_img_res = {
  383. "input_path": batch_data.input_paths[0],
  384. "page_index": batch_data.page_indexes[0],
  385. "doc_preprocessor_res": doc_preprocessor_res,
  386. "layout_det_res": layout_det_res,
  387. "overall_ocr_res": overall_ocr_res,
  388. "table_res_list": table_res_list,
  389. "model_settings": model_settings,
  390. }
  391. yield TableRecognitionResult(single_img_res)