pipeline.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. from typing import Any, Dict, List, Optional, Tuple, Union
  15. import numpy as np
  16. from ....utils import logging
  17. from ...common.batch_sampler import ImageBatchSampler
  18. from ...common.reader import ReadImage
  19. from ...models.formula_recognition.result import (
  20. FormulaRecResult as SingleFormulaRecognitionResult,
  21. )
  22. from ...models.object_detection.result import DetResult
  23. from ...utils.hpi import HPIConfig
  24. from ...utils.pp_option import PaddlePredictorOption
  25. from ..base import BasePipeline
  26. from ..components import CropByBoxes
  27. from .result import FormulaRecognitionResult
  28. class FormulaRecognitionPipeline(BasePipeline):
  29. """Formula Recognition Pipeline"""
  30. entities = ["formula_recognition"]
  31. def __init__(
  32. self,
  33. config: Dict,
  34. device: str = None,
  35. pp_option: PaddlePredictorOption = None,
  36. use_hpip: bool = False,
  37. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  38. ) -> None:
  39. """Initializes the formula recognition pipeline.
  40. Args:
  41. config (Dict): Configuration dictionary containing various settings.
  42. device (str, optional): Device to run the predictions on. Defaults to None.
  43. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  44. use_hpip (bool, optional): Whether to use the high-performance
  45. inference plugin (HPIP) by default. Defaults to False.
  46. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  47. The default high-performance inference configuration dictionary.
  48. Defaults to None.
  49. """
  50. super().__init__(
  51. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  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. layout_kwargs = {}
  71. if (threshold := layout_det_config.get("threshold", None)) is not None:
  72. layout_kwargs["threshold"] = threshold
  73. if (layout_nms := layout_det_config.get("layout_nms", None)) is not None:
  74. layout_kwargs["layout_nms"] = layout_nms
  75. if (
  76. layout_unclip_ratio := layout_det_config.get(
  77. "layout_unclip_ratio", None
  78. )
  79. ) is not None:
  80. layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio
  81. if (
  82. layout_merge_bboxes_mode := layout_det_config.get(
  83. "layout_merge_bboxes_mode", None
  84. )
  85. ) is not None:
  86. layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode
  87. self.layout_det_model = self.create_model(
  88. layout_det_config, **layout_kwargs
  89. )
  90. formula_recognition_config = config.get("SubModules", {}).get(
  91. "FormulaRecognition",
  92. {"model_config_error": "config error for formula_rec_model!"},
  93. )
  94. self.formula_recognition_model = self.create_model(formula_recognition_config)
  95. self._crop_by_boxes = CropByBoxes()
  96. self.batch_sampler = ImageBatchSampler(batch_size=1)
  97. self.img_reader = ReadImage(format="BGR")
  98. def get_model_settings(
  99. self,
  100. use_doc_orientation_classify: Optional[bool],
  101. use_doc_unwarping: Optional[bool],
  102. use_layout_detection: Optional[bool],
  103. ) -> dict:
  104. """
  105. Get the model settings based on the provided parameters or default values.
  106. Args:
  107. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  108. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  109. use_layout_detection (Optional[bool]): Whether to use layout detection.
  110. Returns:
  111. dict: A dictionary containing the model settings.
  112. """
  113. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  114. use_doc_preprocessor = self.use_doc_preprocessor
  115. else:
  116. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  117. use_doc_preprocessor = True
  118. else:
  119. use_doc_preprocessor = False
  120. if use_layout_detection is None:
  121. use_layout_detection = self.use_layout_detection
  122. return dict(
  123. use_doc_preprocessor=use_doc_preprocessor,
  124. use_layout_detection=use_layout_detection,
  125. )
  126. def check_model_settings_valid(
  127. self, model_settings: Dict, layout_det_res: DetResult
  128. ) -> bool:
  129. """
  130. Check if the input parameters are valid based on the initialized models.
  131. Args:
  132. model_settings (Dict): A dictionary containing input parameters.
  133. layout_det_res (DetResult): The layout detection result.
  134. Returns:
  135. bool: True if all required models are initialized according to input parameters, False otherwise.
  136. """
  137. if model_settings["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  138. logging.error(
  139. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  140. )
  141. return False
  142. if model_settings["use_layout_detection"]:
  143. if layout_det_res is not None:
  144. logging.error(
  145. "The layout detection model has already been initialized, please set use_layout_detection=False"
  146. )
  147. return False
  148. if not self.use_layout_detection:
  149. logging.error(
  150. "Set use_layout_detection, but the models for layout detection are not initialized."
  151. )
  152. return False
  153. return True
  154. def predict_single_formula_recognition_res(
  155. self,
  156. image_array: np.ndarray,
  157. ) -> SingleFormulaRecognitionResult:
  158. """
  159. Predict formula recognition results from an image array, layout detection results.
  160. Args:
  161. image_array (np.ndarray): The input image represented as a numpy array.
  162. formula_box (list): The formula box coordinates.
  163. flag_find_nei_text (bool): Whether to find neighboring text.
  164. Returns:
  165. SingleFormulaRecognitionResult: single formula recognition result.
  166. """
  167. formula_recognition_pred = next(self.formula_recognition_model(image_array))
  168. return formula_recognition_pred
  169. def predict(
  170. self,
  171. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  172. use_layout_detection: Optional[bool] = None,
  173. use_doc_orientation_classify: Optional[bool] = None,
  174. use_doc_unwarping: Optional[bool] = None,
  175. layout_det_res: Optional[DetResult] = None,
  176. layout_threshold: Optional[Union[float, dict]] = None,
  177. layout_nms: Optional[bool] = None,
  178. layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
  179. layout_merge_bboxes_mode: Optional[str] = None,
  180. **kwargs,
  181. ) -> FormulaRecognitionResult:
  182. """
  183. This function predicts the layout parsing result for the given input.
  184. Args:
  185. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) of pdf(s) to be processed.
  186. use_layout_detection (Optional[bool]): Whether to use layout detection.
  187. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  188. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  189. layout_det_res (Optional[DetResult]): The layout detection result.
  190. It will be used if it is not None and use_layout_detection is False.
  191. **kwargs: Additional keyword arguments.
  192. Returns:
  193. formulaRecognitionResult: The predicted formula recognition result.
  194. """
  195. model_settings = self.get_model_settings(
  196. use_doc_orientation_classify,
  197. use_doc_unwarping,
  198. use_layout_detection,
  199. )
  200. if not self.check_model_settings_valid(model_settings, layout_det_res):
  201. yield {"error": "the input params for model settings are invalid!"}
  202. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  203. image_array = self.img_reader(batch_data.instances)[0]
  204. if model_settings["use_doc_preprocessor"]:
  205. doc_preprocessor_res = next(
  206. self.doc_preprocessor_pipeline(
  207. image_array,
  208. use_doc_orientation_classify=use_doc_orientation_classify,
  209. use_doc_unwarping=use_doc_unwarping,
  210. )
  211. )
  212. else:
  213. doc_preprocessor_res = {"output_img": image_array}
  214. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  215. formula_res_list = []
  216. formula_region_id = 1
  217. if not model_settings["use_layout_detection"] and layout_det_res is None:
  218. layout_det_res = {}
  219. img_height, img_width = doc_preprocessor_image.shape[:2]
  220. single_formula_rec_res = self.predict_single_formula_recognition_res(
  221. doc_preprocessor_image,
  222. )
  223. single_formula_rec_res["formula_region_id"] = formula_region_id
  224. formula_res_list.append(single_formula_rec_res)
  225. formula_region_id += 1
  226. else:
  227. if model_settings["use_layout_detection"]:
  228. layout_det_res = next(
  229. self.layout_det_model(
  230. doc_preprocessor_image,
  231. threshold=layout_threshold,
  232. layout_nms=layout_nms,
  233. layout_unclip_ratio=layout_unclip_ratio,
  234. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  235. )
  236. )
  237. formula_crop_img = []
  238. for box_info in layout_det_res["boxes"]:
  239. if box_info["label"].lower() in ["formula"]:
  240. crop_img_info = self._crop_by_boxes(
  241. doc_preprocessor_image, [box_info]
  242. )
  243. crop_img_info = crop_img_info[0]
  244. formula_crop_img.append(crop_img_info["img"])
  245. single_formula_rec_res = {}
  246. single_formula_rec_res["formula_region_id"] = formula_region_id
  247. single_formula_rec_res["dt_polys"] = box_info["coordinate"]
  248. formula_res_list.append(single_formula_rec_res)
  249. formula_region_id += 1
  250. for idx, formula_rec_res in enumerate(
  251. self.formula_recognition_model(formula_crop_img)
  252. ):
  253. formula_region_id = formula_res_list[idx]["formula_region_id"]
  254. dt_polys = formula_res_list[idx]["dt_polys"]
  255. formula_rec_res["formula_region_id"] = formula_region_id
  256. formula_rec_res["dt_polys"] = dt_polys
  257. formula_res_list[idx] = formula_rec_res
  258. single_img_res = {
  259. "input_path": batch_data.input_paths[0],
  260. "page_index": batch_data.page_indexes[0],
  261. "layout_det_res": layout_det_res,
  262. "doc_preprocessor_res": doc_preprocessor_res,
  263. "formula_res_list": formula_res_list,
  264. "model_settings": model_settings,
  265. }
  266. yield FormulaRecognitionResult(single_img_res)