table_recognition.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190
  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 numpy as np
  15. from .utils import *
  16. from ..base import BasePipeline
  17. from ..ocr import OCRPipeline
  18. from ....utils import logging
  19. from ...components import CropByBoxes
  20. from ...results import OCRResult, TableResult, StructureTableResult
  21. class _TableRecPipeline(BasePipeline):
  22. """Table Recognition Pipeline"""
  23. def __init__(
  24. self,
  25. predictor_kwargs=None,
  26. ):
  27. super().__init__(predictor_kwargs=predictor_kwargs)
  28. def _build_predictor(
  29. self,
  30. layout_model,
  31. text_det_model,
  32. text_rec_model,
  33. table_model,
  34. ):
  35. self.layout_predictor = self._create(model=layout_model)
  36. self.ocr_pipeline = self._create(
  37. pipeline=OCRPipeline,
  38. text_det_model=text_det_model,
  39. text_rec_model=text_rec_model,
  40. )
  41. self.table_predictor = self._create(model=table_model)
  42. self._crop_by_boxes = CropByBoxes()
  43. self._match = TableMatch(filter_ocr_result=False)
  44. def set_predictor(
  45. self,
  46. layout_batch_size=None,
  47. text_det_batch_size=None,
  48. text_rec_batch_size=None,
  49. table_batch_size=None,
  50. device=None,
  51. ):
  52. if text_det_batch_size and text_det_batch_size > 1:
  53. logging.warning(
  54. f"text det model only support batch_size=1 now,the setting of text_det_batch_size={text_det_batch_size} will not using! "
  55. )
  56. if layout_batch_size:
  57. self.layout_predictor.set_predictor(batch_size=layout_batch_size)
  58. if text_rec_batch_size:
  59. self.ocr_pipeline.text_rec_model.set_predictor(
  60. batch_size=text_rec_batch_size
  61. )
  62. if table_batch_size:
  63. self.table_predictor.set_predictor(batch_size=table_batch_size)
  64. if device:
  65. self.layout_predictor.set_predictor(device=device)
  66. self.ocr_pipeline.text_rec_model.set_predictor(device=device)
  67. self.table_predictor.set_predictor(device=device)
  68. def predict(self, input, **kwargs):
  69. self.set_predictor(**kwargs)
  70. for layout_pred, ocr_pred in zip(
  71. self.layout_predictor(input), self.ocr_pipeline(input)
  72. ):
  73. single_img_res = {
  74. "input_path": "",
  75. "layout_result": {},
  76. "ocr_result": {},
  77. "table_result": [],
  78. }
  79. # update layout result
  80. single_img_res["input_path"] = layout_pred["input_path"]
  81. single_img_res["layout_result"] = layout_pred
  82. ocr_res = ocr_pred
  83. table_subs = []
  84. if len(layout_pred["boxes"]) > 0:
  85. subs_of_img = list(self._crop_by_boxes(layout_pred))
  86. # get cropped images with label "table"
  87. for sub in subs_of_img:
  88. box = sub["box"]
  89. if sub["label"].lower() == "table":
  90. table_subs.append(sub)
  91. _, ocr_res = self.get_related_ocr_result(box, ocr_res)
  92. table_res, all_table_ocr_res = self.get_table_result(table_subs)
  93. for table_ocr_res in all_table_ocr_res:
  94. ocr_res["dt_polys"].extend(table_ocr_res["dt_polys"])
  95. ocr_res["rec_text"].extend(table_ocr_res["rec_text"])
  96. ocr_res["rec_score"].extend(table_ocr_res["rec_score"])
  97. single_img_res["table_result"] = table_res
  98. single_img_res["ocr_result"] = OCRResult(ocr_res)
  99. yield TableResult(single_img_res)
  100. def get_related_ocr_result(self, box, ocr_res):
  101. dt_polys_list = []
  102. rec_text_list = []
  103. score_list = []
  104. unmatched_ocr_res = {"dt_polys": [], "rec_text": [], "rec_score": []}
  105. unmatched_ocr_res["input_path"] = ocr_res["input_path"]
  106. for i, text_box in enumerate(ocr_res["dt_polys"]):
  107. text_box_area = convert_4point2rect(text_box)
  108. if is_inside(text_box_area, box):
  109. dt_polys_list.append(text_box)
  110. rec_text_list.append(ocr_res["rec_text"][i])
  111. score_list.append(ocr_res["rec_score"][i])
  112. else:
  113. unmatched_ocr_res["dt_polys"].append(text_box)
  114. unmatched_ocr_res["rec_text"].append(ocr_res["rec_text"][i])
  115. unmatched_ocr_res["rec_score"].append(ocr_res["rec_score"][i])
  116. return (dt_polys_list, rec_text_list, score_list), unmatched_ocr_res
  117. def get_table_result(self, input_imgs):
  118. table_res_list = []
  119. ocr_res_list = []
  120. table_index = 0
  121. img_list = [img["img"] for img in input_imgs]
  122. for input_img, table_pred, ocr_pred in zip(
  123. input_imgs, self.table_predictor(img_list), self.ocr_pipeline(img_list)
  124. ):
  125. single_table_box = table_pred["bbox"]
  126. ori_x, ori_y, _, _ = input_img["box"]
  127. ori_bbox_list = np.array(
  128. get_ori_coordinate_for_table(ori_x, ori_y, single_table_box),
  129. dtype=np.float32,
  130. )
  131. ori_ocr_bbox_list = np.array(
  132. get_ori_coordinate_for_table(ori_x, ori_y, ocr_pred["dt_polys"]),
  133. dtype=np.float32,
  134. )
  135. html_res = self._match(table_pred, ocr_pred)
  136. ocr_pred["dt_polys"] = ori_ocr_bbox_list
  137. table_res_list.append(
  138. StructureTableResult(
  139. {
  140. "input_path": input_img["input_path"],
  141. "layout_bbox": [int(x) for x in input_img["box"]],
  142. "bbox": ori_bbox_list,
  143. "img_idx": table_index,
  144. "html": html_res,
  145. }
  146. )
  147. )
  148. ocr_res_list.append(ocr_pred)
  149. table_index += 1
  150. return table_res_list, ocr_res_list
  151. class TableRecPipeline(_TableRecPipeline):
  152. """Table Recognition Pipeline"""
  153. entities = "table_recognition"
  154. def __init__(
  155. self,
  156. layout_model,
  157. text_det_model,
  158. text_rec_model,
  159. table_model,
  160. layout_batch_size=1,
  161. text_det_batch_size=1,
  162. text_rec_batch_size=1,
  163. table_batch_size=1,
  164. device=None,
  165. predictor_kwargs=None,
  166. ):
  167. super().__init__(predictor_kwargs=predictor_kwargs)
  168. self._build_predictor(layout_model, text_det_model, text_rec_model, table_model)
  169. self.set_predictor(
  170. layout_batch_size=layout_batch_size,
  171. text_det_batch_size=text_det_batch_size,
  172. text_rec_batch_size=text_rec_batch_size,
  173. table_batch_size=table_batch_size,
  174. device=device,
  175. )