pipeline_v2.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669
  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 sklearn.cluster import KMeans
  19. from ..base import BasePipeline
  20. from ..components import CropByBoxes
  21. from .utils import get_neighbor_boxes_idx
  22. from .table_recognition_post_processing_v2 import get_table_recognition_res
  23. from .result import SingleTableRecognitionResult, TableRecognitionResult
  24. from ....utils import logging
  25. from ...utils.pp_option import PaddlePredictorOption
  26. from ...common.reader import ReadImage
  27. from ...common.batch_sampler import ImageBatchSampler
  28. from ..ocr.result import OCRResult
  29. from ..doc_preprocessor.result import DocPreprocessorResult
  30. from ...models.object_detection.result import DetResult
  31. class TableRecognitionPipelineV2(BasePipeline):
  32. """Table Recognition Pipeline"""
  33. entities = ["table_recognition_v2"]
  34. def __init__(
  35. self,
  36. config: Dict,
  37. device: str = None,
  38. pp_option: PaddlePredictorOption = None,
  39. use_hpip: bool = False,
  40. hpi_params: Optional[Dict[str, Any]] = None,
  41. ) -> None:
  42. """Initializes the layout parsing pipeline.
  43. Args:
  44. config (Dict): Configuration dictionary containing various settings.
  45. device (str, optional): Device to run the predictions on. Defaults to None.
  46. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  47. use_hpip (bool, optional): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
  48. hpi_params (Optional[Dict[str, Any]], optional): HPIP parameters. Defaults to None.
  49. """
  50. super().__init__(
  51. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_params=hpi_params
  52. )
  53. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  54. if self.use_doc_preprocessor:
  55. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  56. "DocPreprocessor",
  57. {
  58. "pipeline_config_error": "config error for doc_preprocessor_pipeline!"
  59. },
  60. )
  61. self.doc_preprocessor_pipeline = self.create_pipeline(
  62. doc_preprocessor_config
  63. )
  64. self.use_layout_detection = config.get("use_layout_detection", True)
  65. if self.use_layout_detection:
  66. layout_det_config = config.get("SubModules", {}).get(
  67. "LayoutDetection",
  68. {"model_config_error": "config error for layout_det_model!"},
  69. )
  70. self.layout_det_model = self.create_model(layout_det_config)
  71. table_cls_config = config.get("SubModules", {}).get(
  72. "TableClassification",
  73. {"model_config_error": "config error for table_classification_model!"},
  74. )
  75. self.table_cls_model = self.create_model(table_cls_config)
  76. wired_table_rec_config = config.get("SubModules", {}).get(
  77. "WiredTableStructureRecognition",
  78. {"model_config_error": "config error for wired_table_structure_model!"},
  79. )
  80. self.wired_table_rec_model = self.create_model(wired_table_rec_config)
  81. wireless_table_rec_config = config.get("SubModules", {}).get(
  82. "WirelessTableStructureRecognition",
  83. {"model_config_error": "config error for wireless_table_structure_model!"},
  84. )
  85. self.wireless_table_rec_model = self.create_model(wireless_table_rec_config)
  86. wired_table_cells_det_config = config.get("SubModules", {}).get(
  87. "WiredTableCellsDetection",
  88. {
  89. "model_config_error": "config error for wired_table_cells_detection_model!"
  90. },
  91. )
  92. self.wired_table_cells_detection_model = self.create_model(
  93. wired_table_cells_det_config
  94. )
  95. wireless_table_cells_det_config = config.get("SubModules", {}).get(
  96. "WirelessTableCellsDetection",
  97. {
  98. "model_config_error": "config error for wireless_table_cells_detection_model!"
  99. },
  100. )
  101. self.wireless_table_cells_detection_model = self.create_model(
  102. wireless_table_cells_det_config
  103. )
  104. self.use_ocr_model = config.get("use_ocr_model", True)
  105. if self.use_ocr_model:
  106. general_ocr_config = config.get("SubPipelines", {}).get(
  107. "GeneralOCR",
  108. {"pipeline_config_error": "config error for general_ocr_pipeline!"},
  109. )
  110. self.general_ocr_pipeline = self.create_pipeline(general_ocr_config)
  111. self._crop_by_boxes = CropByBoxes()
  112. self.batch_sampler = ImageBatchSampler(batch_size=1)
  113. self.img_reader = ReadImage(format="BGR")
  114. def get_model_settings(
  115. self,
  116. use_doc_orientation_classify: Optional[bool],
  117. use_doc_unwarping: Optional[bool],
  118. use_layout_detection: Optional[bool],
  119. use_ocr_model: Optional[bool],
  120. ) -> dict:
  121. """
  122. Get the model settings based on the provided parameters or default values.
  123. Args:
  124. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  125. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  126. use_layout_detection (Optional[bool]): Whether to use layout detection.
  127. use_ocr_model (Optional[bool]): Whether to use OCR model.
  128. Returns:
  129. dict: A dictionary containing the model settings.
  130. """
  131. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  132. use_doc_preprocessor = self.use_doc_preprocessor
  133. else:
  134. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  135. use_doc_preprocessor = True
  136. else:
  137. use_doc_preprocessor = False
  138. if use_layout_detection is None:
  139. use_layout_detection = self.use_layout_detection
  140. if use_ocr_model is None:
  141. use_ocr_model = self.use_ocr_model
  142. return dict(
  143. use_doc_preprocessor=use_doc_preprocessor,
  144. use_layout_detection=use_layout_detection,
  145. use_ocr_model=use_ocr_model,
  146. )
  147. def check_model_settings_valid(
  148. self,
  149. model_settings: Dict,
  150. overall_ocr_res: OCRResult,
  151. layout_det_res: DetResult,
  152. ) -> bool:
  153. """
  154. Check if the input parameters are valid based on the initialized models.
  155. Args:
  156. model_settings (Dict): A dictionary containing input parameters.
  157. overall_ocr_res (OCRResult): Overall OCR result obtained after running the OCR pipeline.
  158. The overall OCR result with convert_points_to_boxes information.
  159. layout_det_res (DetResult): The layout detection result.
  160. Returns:
  161. bool: True if all required models are initialized according to input parameters, False otherwise.
  162. """
  163. if model_settings["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  164. logging.error(
  165. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  166. )
  167. return False
  168. if model_settings["use_layout_detection"]:
  169. if layout_det_res is not None:
  170. logging.error(
  171. "The layout detection model has already been initialized, please set use_layout_detection=False"
  172. )
  173. return False
  174. if not self.use_layout_detection:
  175. logging.error(
  176. "Set use_layout_detection, but the models for layout detection are not initialized."
  177. )
  178. return False
  179. if model_settings["use_ocr_model"]:
  180. if overall_ocr_res is not None:
  181. logging.error(
  182. "The OCR models have already been initialized, please set use_ocr_model=False"
  183. )
  184. return False
  185. if not self.use_ocr_model:
  186. logging.error(
  187. "Set use_ocr_model, but the models for OCR are not initialized."
  188. )
  189. return False
  190. else:
  191. if overall_ocr_res is None:
  192. logging.error("Set use_ocr_model=False, but no OCR results were found.")
  193. return False
  194. return True
  195. def predict_doc_preprocessor_res(
  196. self, image_array: np.ndarray, input_params: dict
  197. ) -> Tuple[DocPreprocessorResult, np.ndarray]:
  198. """
  199. Preprocess the document image based on input parameters.
  200. Args:
  201. image_array (np.ndarray): The input image array.
  202. input_params (dict): Dictionary containing preprocessing parameters.
  203. Returns:
  204. tuple[DocPreprocessorResult, np.ndarray]: A tuple containing the preprocessing
  205. result dictionary and the processed image array.
  206. """
  207. if input_params["use_doc_preprocessor"]:
  208. use_doc_orientation_classify = input_params["use_doc_orientation_classify"]
  209. use_doc_unwarping = input_params["use_doc_unwarping"]
  210. doc_preprocessor_res = next(
  211. self.doc_preprocessor_pipeline(
  212. image_array,
  213. use_doc_orientation_classify=use_doc_orientation_classify,
  214. use_doc_unwarping=use_doc_unwarping,
  215. )
  216. )
  217. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  218. else:
  219. doc_preprocessor_res = {}
  220. doc_preprocessor_image = image_array
  221. return doc_preprocessor_res, doc_preprocessor_image
  222. def extract_results(self, pred, task):
  223. if task == "cls":
  224. return pred["label_names"][np.argmax(pred["scores"])]
  225. elif task == "det":
  226. threshold = 0.0
  227. result = []
  228. cell_score = []
  229. if "boxes" in pred and isinstance(pred["boxes"], list):
  230. for box in pred["boxes"]:
  231. if isinstance(box, dict) and "score" in box and "coordinate" in box:
  232. score = box["score"]
  233. coordinate = box["coordinate"]
  234. if isinstance(score, float) and score > threshold:
  235. result.append(coordinate)
  236. cell_score.append(score)
  237. return result, cell_score
  238. elif task == "table_stru":
  239. return pred["structure"]
  240. else:
  241. return None
  242. def cells_det_results_nms(self, cells_det_results, cells_det_scores, cells_det_threshold=0.3):
  243. """
  244. Apply Non-Maximum Suppression (NMS) on detection results to remove redundant overlapping bounding boxes.
  245. Args:
  246. cells_det_results (list): List of bounding boxes, each box is in format [x1, y1, x2, y2].
  247. cells_det_scores (list): List of confidence scores corresponding to the bounding boxes.
  248. cells_det_threshold (float): IoU threshold for suppression. Boxes with IoU greater than this threshold
  249. will be suppressed. Default is 0.5.
  250. Returns:
  251. Tuple[list, list]: A tuple containing the list of bounding boxes and confidence scores after NMS,
  252. while maintaining one-to-one correspondence.
  253. """
  254. # Convert lists to numpy arrays for efficient computation
  255. boxes = np.array(cells_det_results)
  256. scores = np.array(cells_det_scores)
  257. # Initialize list for picked indices
  258. picked_indices = []
  259. # Get coordinates of bounding boxes
  260. x1 = boxes[:, 0]
  261. y1 = boxes[:, 1]
  262. x2 = boxes[:, 2]
  263. y2 = boxes[:, 3]
  264. # Compute the area of the bounding boxes
  265. areas = (x2 - x1) * (y2 - y1)
  266. # Sort the bounding boxes by the confidence scores in descending order
  267. order = scores.argsort()[::-1]
  268. # Process the boxes
  269. while order.size > 0:
  270. # Index of the current highest score box
  271. i = order[0]
  272. picked_indices.append(i)
  273. # Compute IoU between the highest score box and the rest
  274. xx1 = np.maximum(x1[i], x1[order[1:]])
  275. yy1 = np.maximum(y1[i], y1[order[1:]])
  276. xx2 = np.minimum(x2[i], x2[order[1:]])
  277. yy2 = np.minimum(y2[i], y2[order[1:]])
  278. # Compute the width and height of the overlapping area
  279. w = np.maximum(0.0, xx2 - xx1)
  280. h = np.maximum(0.0, yy2 - yy1)
  281. # Compute the ratio of overlap (IoU)
  282. inter = w * h
  283. ovr = inter / (areas[i] + areas[order[1:]] - inter)
  284. # Indices of boxes with IoU less than threshold
  285. inds = np.where(ovr <= cells_det_threshold)[0]
  286. # Update order, only keep boxes with IoU less than threshold
  287. order = order[inds + 1] # inds shifted by 1 because order[0] is the current box
  288. # Select the boxes and scores based on picked indices
  289. final_boxes = boxes[picked_indices].tolist()
  290. final_scores = scores[picked_indices].tolist()
  291. return final_boxes, final_scores
  292. def get_region_ocr_det_boxes(self, ocr_det_boxes, table_box):
  293. """Adjust the coordinates of ocr_det_boxes that are fully inside table_box relative to table_box.
  294. Args:
  295. ocr_det_boxes (list of list): List of bounding boxes [x1, y1, x2, y2] in the original image.
  296. table_box (list): Bounding box [x1, y1, x2, y2] of the target region in the original image.
  297. Returns:
  298. list of list: List of adjusted bounding boxes relative to table_box, for boxes fully inside table_box.
  299. """
  300. tol = 5
  301. # Extract coordinates from table_box
  302. x_min_t, y_min_t, x_max_t, y_max_t = table_box
  303. adjusted_boxes = []
  304. for box in ocr_det_boxes:
  305. x_min_b, y_min_b, x_max_b, y_max_b = box
  306. # Check if the box is fully inside table_box
  307. if (x_min_b+tol >= x_min_t and y_min_b+tol >= y_min_t and
  308. x_max_b+tol <= x_max_t and y_max_b+tol <= y_max_t):
  309. # Adjust the coordinates to be relative to table_box
  310. adjusted_box = [
  311. x_min_b - x_min_t, # Adjust x1
  312. y_min_b - y_min_t, # Adjust y1
  313. x_max_b - x_min_t, # Adjust x2
  314. y_max_b - y_min_t # Adjust y2
  315. ]
  316. adjusted_boxes.append(adjusted_box)
  317. # Discard boxes not fully inside table_box
  318. return adjusted_boxes
  319. def cells_det_results_reprocessing(self, cells_det_results, cells_det_scores, ocr_det_results, html_pred_boxes_nums):
  320. """
  321. Process and filter cells_det_results based on ocr_det_results and html_pred_boxes_nums.
  322. Args:
  323. cells_det_results (List[List[float]]): List of detected cell rectangles [[x1, y1, x2, y2], ...].
  324. cells_det_scores (List[float]): List of confidence scores for each rectangle in cells_det_results.
  325. ocr_det_results (List[List[float]]): List of OCR detected rectangles [[x1, y1, x2, y2], ...].
  326. html_pred_boxes_nums (int): The desired number of rectangles in the final output.
  327. Returns:
  328. List[List[float]]: The processed list of rectangles.
  329. """
  330. # Function to compute IoU between two rectangles
  331. def compute_iou(box1, box2):
  332. """
  333. Compute the Intersection over Union (IoU) between two rectangles.
  334. Args:
  335. box1 (array-like): [x1, y1, x2, y2] of the first rectangle.
  336. box2 (array-like): [x1, y1, x2, y2] of the second rectangle.
  337. Returns:
  338. float: The IoU between the two rectangles.
  339. """
  340. # Determine the coordinates of the intersection rectangle
  341. x_left = max(box1[0], box2[0])
  342. y_top = max(box1[1], box2[1])
  343. x_right = min(box1[2], box2[2])
  344. y_bottom = min(box1[3], box2[3])
  345. if x_right <= x_left or y_bottom <= y_top:
  346. return 0.0
  347. # Calculate the area of intersection rectangle
  348. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  349. # Calculate the area of both rectangles
  350. box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
  351. box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
  352. # Calculate the IoU
  353. iou = intersection_area / float(box1_area + box2_area - intersection_area)
  354. return iou
  355. # Function to combine rectangles into N rectangles
  356. def combine_rectangles(rectangles, N):
  357. """
  358. Combine rectangles into N rectangles based on geometric proximity.
  359. Args:
  360. rectangles (list of list of int): A list of rectangles, each represented by [x1, y1, x2, y2].
  361. N (int): The desired number of combined rectangles.
  362. Returns:
  363. list of list of int: A list of N combined rectangles.
  364. """
  365. # Number of input rectangles
  366. num_rects = len(rectangles)
  367. # If N is greater than or equal to the number of rectangles, return the original rectangles
  368. if N >= num_rects:
  369. return rectangles
  370. # Compute the center points of the rectangles
  371. centers = np.array([
  372. [
  373. (rect[0] + rect[2]) / 2, # Center x-coordinate
  374. (rect[1] + rect[3]) / 2 # Center y-coordinate
  375. ]
  376. for rect in rectangles
  377. ])
  378. # Perform KMeans clustering on the center points to group them into N clusters
  379. kmeans = KMeans(n_clusters=N, random_state=0, n_init='auto')
  380. labels = kmeans.fit_predict(centers)
  381. # Initialize a list to store the combined rectangles
  382. combined_rectangles = []
  383. # For each cluster, compute the minimal bounding rectangle that covers all rectangles in the cluster
  384. for i in range(N):
  385. # Get the indices of rectangles that belong to cluster i
  386. indices = np.where(labels == i)[0]
  387. if len(indices) == 0:
  388. # If no rectangles in this cluster, skip it
  389. continue
  390. # Extract the rectangles in cluster i
  391. cluster_rects = np.array([rectangles[idx] for idx in indices])
  392. # Compute the minimal x1, y1 (top-left corner) and maximal x2, y2 (bottom-right corner)
  393. x1_min = np.min(cluster_rects[:, 0])
  394. y1_min = np.min(cluster_rects[:, 1])
  395. x2_max = np.max(cluster_rects[:, 2])
  396. y2_max = np.max(cluster_rects[:, 3])
  397. # Append the combined rectangle to the list
  398. combined_rectangles.append([x1_min, y1_min, x2_max, y2_max])
  399. return combined_rectangles
  400. # Ensure that the inputs are numpy arrays for efficient computation
  401. cells_det_results = np.array(cells_det_results)
  402. cells_det_scores = np.array(cells_det_scores)
  403. ocr_det_results = np.array(ocr_det_results)
  404. if len(cells_det_results) == html_pred_boxes_nums:
  405. return cells_det_results
  406. # Step 1: If cells_det_results has more rectangles than html_pred_boxes_nums
  407. elif len(cells_det_results) > html_pred_boxes_nums:
  408. return combine_rectangles(cells_det_results, html_pred_boxes_nums)
  409. else:
  410. # return cells_det_results
  411. # Threshold for IoU
  412. iou_threshold = 0.1
  413. # List to store ocr_miss_boxes
  414. ocr_miss_boxes = []
  415. # For each rectangle in ocr_det_results
  416. for ocr_rect in ocr_det_results:
  417. # Flag to indicate if ocr_rect has IoU >= threshold with any cell_rect
  418. has_large_iou = False
  419. # For each rectangle in cells_det_results
  420. for cell_rect in cells_det_results:
  421. # Compute IoU
  422. iou = compute_iou(ocr_rect, cell_rect)
  423. if iou >= iou_threshold:
  424. has_large_iou = True
  425. break
  426. if not has_large_iou:
  427. ocr_miss_boxes.append(ocr_rect)
  428. # If no ocr_miss_boxes, return cells_det_results
  429. if len(ocr_miss_boxes) == 0:
  430. return cells_det_results.tolist()
  431. else:
  432. # Need to combine ocr_miss_boxes into N rectangles
  433. N = html_pred_boxes_nums - len(cells_det_results)
  434. if len(ocr_miss_boxes) == N:
  435. return cells_det_results.tolist() + ocr_miss_boxes
  436. else:
  437. # Combine ocr_miss_boxes into N rectangles
  438. ocr_supp_boxes = combine_rectangles(ocr_miss_boxes, N)
  439. # Combine cells_det_results and ocr_supp_boxes
  440. final_results = np.concatenate((cells_det_results, ocr_supp_boxes), axis=0)
  441. return final_results.tolist()
  442. def predict_single_table_recognition_res(
  443. self,
  444. image_array: np.ndarray,
  445. overall_ocr_res: OCRResult,
  446. table_box: list,
  447. flag_find_nei_text: bool = True,
  448. ) -> SingleTableRecognitionResult:
  449. """
  450. Predict table recognition results from an image array, layout detection results, and OCR results.
  451. Args:
  452. image_array (np.ndarray): The input image represented as a numpy array.
  453. overall_ocr_res (OCRResult): Overall OCR result obtained after running the OCR pipeline.
  454. The overall OCR results containing text recognition information.
  455. table_box (list): The table box coordinates.
  456. flag_find_nei_text (bool): Whether to find neighboring text.
  457. Returns:
  458. SingleTableRecognitionResult: single table recognition result.
  459. """
  460. table_cls_pred = next(self.table_cls_model(image_array))
  461. table_cls_result = self.extract_results(table_cls_pred, "cls")
  462. if table_cls_result == "wired_table":
  463. table_structure_pred = next(self.wired_table_rec_model(image_array))
  464. table_cells_pred = next(
  465. self.wired_table_cells_detection_model(image_array, threshold=0.3)
  466. ) # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
  467. # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
  468. elif table_cls_result == "wireless_table":
  469. table_structure_pred = next(self.wireless_table_rec_model(image_array))
  470. table_cells_pred = next(
  471. self.wireless_table_cells_detection_model(image_array, threshold=0.3)
  472. ) # Setting the threshold to 0.3 can improve the accuracy of table cells detection.
  473. # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
  474. table_structure_result = self.extract_results(
  475. table_structure_pred, "table_stru"
  476. )
  477. table_cells_result, table_cells_score = self.extract_results(table_cells_pred, "det")
  478. table_cells_result, table_cells_score = self.cells_det_results_nms(table_cells_result, table_cells_score)
  479. ocr_det_boxes = self.get_region_ocr_det_boxes(overall_ocr_res["rec_boxes"].tolist(), table_box)
  480. table_cells_result = self.cells_det_results_reprocessing(
  481. table_cells_result, table_cells_score, ocr_det_boxes, len(table_structure_pred['bbox'])
  482. )
  483. single_table_recognition_res = get_table_recognition_res(
  484. table_box, table_structure_result, table_cells_result, overall_ocr_res
  485. )
  486. neighbor_text = ""
  487. if flag_find_nei_text:
  488. match_idx_list = get_neighbor_boxes_idx(
  489. overall_ocr_res["rec_boxes"], table_box
  490. )
  491. if len(match_idx_list) > 0:
  492. for idx in match_idx_list:
  493. neighbor_text += overall_ocr_res["rec_texts"][idx] + "; "
  494. single_table_recognition_res["neighbor_texts"] = neighbor_text
  495. return single_table_recognition_res
  496. def predict(
  497. self,
  498. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  499. use_doc_orientation_classify: Optional[bool] = None,
  500. use_doc_unwarping: Optional[bool] = None,
  501. use_layout_detection: Optional[bool] = None,
  502. use_ocr_model: Optional[bool] = None,
  503. overall_ocr_res: Optional[OCRResult] = None,
  504. layout_det_res: Optional[DetResult] = None,
  505. text_det_limit_side_len: Optional[int] = None,
  506. text_det_limit_type: Optional[str] = None,
  507. text_det_thresh: Optional[float] = None,
  508. text_det_box_thresh: Optional[float] = None,
  509. text_det_unclip_ratio: Optional[float] = None,
  510. text_rec_score_thresh: Optional[float] = None,
  511. **kwargs,
  512. ) -> TableRecognitionResult:
  513. """
  514. This function predicts the layout parsing result for the given input.
  515. Args:
  516. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) of pdf(s) to be processed.
  517. use_layout_detection (bool): Whether to use layout detection.
  518. use_doc_orientation_classify (bool): Whether to use document orientation classification.
  519. use_doc_unwarping (bool): Whether to use document unwarping.
  520. overall_ocr_res (OCRResult): The overall OCR result with convert_points_to_boxes information.
  521. It will be used if it is not None and use_ocr_model is False.
  522. layout_det_res (DetResult): The layout detection result.
  523. It will be used if it is not None and use_layout_detection is False.
  524. **kwargs: Additional keyword arguments.
  525. Returns:
  526. TableRecognitionResult: The predicted table recognition result.
  527. """
  528. model_settings = self.get_model_settings(
  529. use_doc_orientation_classify,
  530. use_doc_unwarping,
  531. use_layout_detection,
  532. use_ocr_model,
  533. )
  534. if not self.check_model_settings_valid(
  535. model_settings, overall_ocr_res, layout_det_res
  536. ):
  537. yield {"error": "the input params for model settings are invalid!"}
  538. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  539. image_array = self.img_reader(batch_data.instances)[0]
  540. if model_settings["use_doc_preprocessor"]:
  541. doc_preprocessor_res = next(
  542. self.doc_preprocessor_pipeline(
  543. image_array,
  544. use_doc_orientation_classify=use_doc_orientation_classify,
  545. use_doc_unwarping=use_doc_unwarping,
  546. )
  547. )
  548. else:
  549. doc_preprocessor_res = {"output_img": image_array}
  550. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  551. if model_settings["use_ocr_model"]:
  552. overall_ocr_res = next(
  553. self.general_ocr_pipeline(
  554. doc_preprocessor_image,
  555. text_det_limit_side_len=text_det_limit_side_len,
  556. text_det_limit_type=text_det_limit_type,
  557. text_det_thresh=text_det_thresh,
  558. text_det_box_thresh=text_det_box_thresh,
  559. text_det_unclip_ratio=text_det_unclip_ratio,
  560. text_rec_score_thresh=text_rec_score_thresh,
  561. )
  562. )
  563. table_res_list = []
  564. table_region_id = 1
  565. if not model_settings["use_layout_detection"] and layout_det_res is None:
  566. layout_det_res = {}
  567. img_height, img_width = doc_preprocessor_image.shape[:2]
  568. table_box = [0, 0, img_width - 1, img_height - 1]
  569. single_table_rec_res = self.predict_single_table_recognition_res(
  570. doc_preprocessor_image,
  571. overall_ocr_res,
  572. table_box,
  573. flag_find_nei_text=False,
  574. )
  575. single_table_rec_res["table_region_id"] = table_region_id
  576. table_res_list.append(single_table_rec_res)
  577. table_region_id += 1
  578. else:
  579. if model_settings["use_layout_detection"]:
  580. layout_det_res = next(self.layout_det_model(doc_preprocessor_image))
  581. for box_info in layout_det_res["boxes"]:
  582. if box_info["label"].lower() in ["table"]:
  583. crop_img_info = self._crop_by_boxes(image_array, [box_info])
  584. crop_img_info = crop_img_info[0]
  585. table_box = crop_img_info["box"]
  586. single_table_rec_res = (
  587. self.predict_single_table_recognition_res(
  588. crop_img_info["img"], overall_ocr_res, table_box
  589. )
  590. )
  591. single_table_rec_res["table_region_id"] = table_region_id
  592. table_res_list.append(single_table_rec_res)
  593. table_region_id += 1
  594. single_img_res = {
  595. "input_path": batch_data.input_paths[0],
  596. "page_index": batch_data.page_indexes[0],
  597. "doc_preprocessor_res": doc_preprocessor_res,
  598. "layout_det_res": layout_det_res,
  599. "overall_ocr_res": overall_ocr_res,
  600. "table_res_list": table_res_list,
  601. "model_settings": model_settings,
  602. }
  603. yield TableRecognitionResult(single_img_res)