pipeline.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404
  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, 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 ...utils.hpi import HPIConfig
  20. from ...utils.pp_option import PaddlePredictorOption
  21. from ..base import BasePipeline
  22. from ..components import (
  23. CropByPolys,
  24. SortPolyBoxes,
  25. SortQuadBoxes,
  26. convert_points_to_boxes,
  27. rotate_image,
  28. )
  29. from .result import OCRResult
  30. class OCRPipeline(BasePipeline):
  31. """OCR Pipeline"""
  32. entities = "OCR"
  33. def __init__(
  34. self,
  35. config: Dict,
  36. device: Optional[str] = None,
  37. pp_option: Optional[PaddlePredictorOption] = None,
  38. use_hpip: bool = False,
  39. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  40. ) -> None:
  41. """
  42. Initializes the class with given configurations and options.
  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 the high-performance
  48. inference plugin (HPIP) by default. Defaults to False.
  49. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  50. The default high-performance inference configuration dictionary.
  51. Defaults to None.
  52. """
  53. super().__init__(
  54. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  55. )
  56. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  57. if self.use_doc_preprocessor:
  58. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  59. "DocPreprocessor",
  60. {
  61. "pipeline_config_error": "config error for doc_preprocessor_pipeline!"
  62. },
  63. )
  64. self.doc_preprocessor_pipeline = self.create_pipeline(
  65. doc_preprocessor_config
  66. )
  67. self.use_textline_orientation = config.get("use_textline_orientation", True)
  68. if self.use_textline_orientation:
  69. textline_orientation_config = config.get("SubModules", {}).get(
  70. "TextLineOrientation",
  71. {"model_config_error": "config error for textline_orientation_model!"},
  72. )
  73. self.textline_orientation_model = self.create_model(
  74. textline_orientation_config
  75. )
  76. text_det_config = config.get("SubModules", {}).get(
  77. "TextDetection", {"model_config_error": "config error for text_det_model!"}
  78. )
  79. self.text_type = config["text_type"]
  80. if self.text_type == "general":
  81. self.text_det_limit_side_len = text_det_config.get("limit_side_len", 960)
  82. self.text_det_limit_type = text_det_config.get("limit_type", "max")
  83. self.text_det_thresh = text_det_config.get("thresh", 0.3)
  84. self.text_det_box_thresh = text_det_config.get("box_thresh", 0.6)
  85. self.input_shape = text_det_config.get("input_shape", None)
  86. self.text_det_unclip_ratio = text_det_config.get("unclip_ratio", 2.0)
  87. self._sort_boxes = SortQuadBoxes()
  88. self._crop_by_polys = CropByPolys(det_box_type="quad")
  89. elif self.text_type == "seal":
  90. self.text_det_limit_side_len = text_det_config.get("limit_side_len", 736)
  91. self.text_det_limit_type = text_det_config.get("limit_type", "min")
  92. self.text_det_thresh = text_det_config.get("thresh", 0.2)
  93. self.text_det_box_thresh = text_det_config.get("box_thresh", 0.6)
  94. self.text_det_unclip_ratio = text_det_config.get("unclip_ratio", 0.5)
  95. self.input_shape = text_det_config.get("input_shape", None)
  96. self._sort_boxes = SortPolyBoxes()
  97. self._crop_by_polys = CropByPolys(det_box_type="poly")
  98. else:
  99. raise ValueError("Unsupported text type {}".format(self.text_type))
  100. self.text_det_model = self.create_model(
  101. text_det_config,
  102. limit_side_len=self.text_det_limit_side_len,
  103. limit_type=self.text_det_limit_type,
  104. thresh=self.text_det_thresh,
  105. box_thresh=self.text_det_box_thresh,
  106. unclip_ratio=self.text_det_unclip_ratio,
  107. input_shape=self.input_shape,
  108. )
  109. text_rec_config = config.get("SubModules", {}).get(
  110. "TextRecognition",
  111. {"model_config_error": "config error for text_rec_model!"},
  112. )
  113. self.text_rec_score_thresh = text_rec_config.get("score_thresh", 0)
  114. self.input_shape = text_rec_config.get("input_shape", None)
  115. self.text_rec_model = self.create_model(
  116. text_rec_config, input_shape=self.input_shape
  117. )
  118. self.batch_sampler = ImageBatchSampler(batch_size=1)
  119. self.img_reader = ReadImage(format="BGR")
  120. def rotate_image(
  121. self, image_array_list: List[np.ndarray], rotate_angle_list: List[int]
  122. ) -> List[np.ndarray]:
  123. """
  124. Rotate the given image arrays by their corresponding angles.
  125. 0 corresponds to 0 degrees, 1 corresponds to 180 degrees.
  126. Args:
  127. image_array_list (List[np.ndarray]): A list of input image arrays to be rotated.
  128. rotate_angle_list (List[int]): A list of rotation indicators (0 or 1).
  129. 0 means rotate by 0 degrees
  130. 1 means rotate by 180 degrees
  131. Returns:
  132. List[np.ndarray]: A list of rotated image arrays.
  133. Raises:
  134. AssertionError: If any rotate_angle is not 0 or 1.
  135. AssertionError: If the lengths of input lists don't match.
  136. """
  137. assert len(image_array_list) == len(
  138. rotate_angle_list
  139. ), f"Length of image_array_list ({len(image_array_list)}) must match length of rotate_angle_list ({len(rotate_angle_list)})"
  140. for angle in rotate_angle_list:
  141. assert angle in [0, 1], f"rotate_angle must be 0 or 1, now it's {angle}"
  142. rotated_images = []
  143. for image_array, rotate_indicator in zip(image_array_list, rotate_angle_list):
  144. # Convert 0/1 indicator to actual rotation angle
  145. rotate_angle = rotate_indicator * 180
  146. rotated_image = rotate_image(image_array, rotate_angle)
  147. rotated_images.append(rotated_image)
  148. return rotated_images
  149. def check_model_settings_valid(self, model_settings: Dict) -> bool:
  150. """
  151. Check if the input parameters are valid based on the initialized models.
  152. Args:
  153. model_info_params(Dict): A dictionary containing input parameters.
  154. Returns:
  155. bool: True if all required models are initialized according to input parameters, False otherwise.
  156. """
  157. if model_settings["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  158. logging.error(
  159. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  160. )
  161. return False
  162. if (
  163. model_settings["use_textline_orientation"]
  164. and not self.use_textline_orientation
  165. ):
  166. logging.error(
  167. "Set use_textline_orientation, but the models for use_textline_orientation are not initialized."
  168. )
  169. return False
  170. return True
  171. def get_model_settings(
  172. self,
  173. use_doc_orientation_classify: Optional[bool],
  174. use_doc_unwarping: Optional[bool],
  175. use_textline_orientation: Optional[bool],
  176. ) -> dict:
  177. """
  178. Get the model settings based on the provided parameters or default values.
  179. Args:
  180. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  181. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  182. use_textline_orientation (Optional[bool]): Whether to use textline orientation.
  183. Returns:
  184. dict: A dictionary containing the model settings.
  185. """
  186. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  187. use_doc_preprocessor = self.use_doc_preprocessor
  188. else:
  189. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  190. use_doc_preprocessor = True
  191. else:
  192. use_doc_preprocessor = False
  193. if use_textline_orientation is None:
  194. use_textline_orientation = self.use_textline_orientation
  195. return dict(
  196. use_doc_preprocessor=use_doc_preprocessor,
  197. use_textline_orientation=use_textline_orientation,
  198. )
  199. def get_text_det_params(
  200. self,
  201. text_det_limit_side_len: Optional[int] = None,
  202. text_det_limit_type: Optional[str] = None,
  203. text_det_thresh: Optional[float] = None,
  204. text_det_box_thresh: Optional[float] = None,
  205. text_det_unclip_ratio: Optional[float] = None,
  206. ) -> dict:
  207. """
  208. Get text detection parameters.
  209. If a parameter is None, its default value from the instance will be used.
  210. Args:
  211. text_det_limit_side_len (Optional[int]): The maximum side length of the text box.
  212. text_det_limit_type (Optional[str]): The type of limit to apply to the text box.
  213. text_det_thresh (Optional[float]): The threshold for text detection.
  214. text_det_box_thresh (Optional[float]): The threshold for the bounding box.
  215. text_det_unclip_ratio (Optional[float]): The ratio for unclipping the text box.
  216. Returns:
  217. dict: A dictionary containing the text detection parameters.
  218. """
  219. if text_det_limit_side_len is None:
  220. text_det_limit_side_len = self.text_det_limit_side_len
  221. if text_det_limit_type is None:
  222. text_det_limit_type = self.text_det_limit_type
  223. if text_det_thresh is None:
  224. text_det_thresh = self.text_det_thresh
  225. if text_det_box_thresh is None:
  226. text_det_box_thresh = self.text_det_box_thresh
  227. if text_det_unclip_ratio is None:
  228. text_det_unclip_ratio = self.text_det_unclip_ratio
  229. return dict(
  230. limit_side_len=text_det_limit_side_len,
  231. limit_type=text_det_limit_type,
  232. thresh=text_det_thresh,
  233. box_thresh=text_det_box_thresh,
  234. unclip_ratio=text_det_unclip_ratio,
  235. )
  236. def predict(
  237. self,
  238. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  239. use_doc_orientation_classify: Optional[bool] = None,
  240. use_doc_unwarping: Optional[bool] = None,
  241. use_textline_orientation: Optional[bool] = None,
  242. text_det_limit_side_len: Optional[int] = None,
  243. text_det_limit_type: Optional[str] = None,
  244. text_det_thresh: Optional[float] = None,
  245. text_det_box_thresh: Optional[float] = None,
  246. text_det_unclip_ratio: Optional[float] = None,
  247. text_rec_score_thresh: Optional[float] = None,
  248. ) -> OCRResult:
  249. """
  250. Predict OCR results based on input images or arrays with optional preprocessing steps.
  251. Args:
  252. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): Input image of pdf path(s) or numpy array(s).
  253. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  254. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  255. use_textline_orientation (Optional[bool]): Whether to use textline orientation prediction.
  256. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  257. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  258. text_det_thresh (Optional[float]): Threshold for text detection.
  259. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  260. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  261. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  262. Returns:
  263. OCRResult: Generator yielding OCR results for each input image.
  264. """
  265. model_settings = self.get_model_settings(
  266. use_doc_orientation_classify, use_doc_unwarping, use_textline_orientation
  267. )
  268. if not self.check_model_settings_valid(model_settings):
  269. yield {"error": "the input params for model settings are invalid!"}
  270. text_det_params = self.get_text_det_params(
  271. text_det_limit_side_len,
  272. text_det_limit_type,
  273. text_det_thresh,
  274. text_det_box_thresh,
  275. text_det_unclip_ratio,
  276. )
  277. if text_rec_score_thresh is None:
  278. text_rec_score_thresh = self.text_rec_score_thresh
  279. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  280. image_array = self.img_reader(batch_data.instances)[0]
  281. if model_settings["use_doc_preprocessor"]:
  282. doc_preprocessor_res = next(
  283. self.doc_preprocessor_pipeline(
  284. image_array,
  285. use_doc_orientation_classify=use_doc_orientation_classify,
  286. use_doc_unwarping=use_doc_unwarping,
  287. )
  288. )
  289. else:
  290. doc_preprocessor_res = {"output_img": image_array}
  291. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  292. det_res = next(
  293. self.text_det_model(doc_preprocessor_image, **text_det_params)
  294. )
  295. dt_polys = det_res["dt_polys"]
  296. det_res["dt_scores"]
  297. dt_polys = self._sort_boxes(dt_polys)
  298. single_img_res = {
  299. "input_path": batch_data.input_paths[0],
  300. "page_index": batch_data.page_indexes[0],
  301. "doc_preprocessor_res": doc_preprocessor_res,
  302. "dt_polys": dt_polys,
  303. "model_settings": model_settings,
  304. "text_det_params": text_det_params,
  305. "text_type": self.text_type,
  306. "text_rec_score_thresh": text_rec_score_thresh,
  307. }
  308. single_img_res["rec_texts"] = []
  309. single_img_res["rec_scores"] = []
  310. single_img_res["rec_polys"] = []
  311. if len(dt_polys) > 0:
  312. all_subs_of_img = list(
  313. self._crop_by_polys(doc_preprocessor_image, dt_polys)
  314. )
  315. # use textline orientation model
  316. if model_settings["use_textline_orientation"]:
  317. angles = [
  318. int(textline_angle_info["class_ids"][0])
  319. for textline_angle_info in self.textline_orientation_model(
  320. all_subs_of_img
  321. )
  322. ]
  323. all_subs_of_img = self.rotate_image(all_subs_of_img, angles)
  324. else:
  325. angles = [-1] * len(all_subs_of_img)
  326. single_img_res["textline_orientation_angles"] = angles
  327. sub_img_info_list = [
  328. {
  329. "sub_img_id": img_id,
  330. "sub_img_ratio": sub_img.shape[1] / float(sub_img.shape[0]),
  331. }
  332. for img_id, sub_img in enumerate(all_subs_of_img)
  333. ]
  334. sorted_subs_info = sorted(
  335. sub_img_info_list, key=lambda x: x["sub_img_ratio"]
  336. )
  337. sorted_subs_of_img = [
  338. all_subs_of_img[x["sub_img_id"]] for x in sorted_subs_info
  339. ]
  340. for idx, rec_res in enumerate(self.text_rec_model(sorted_subs_of_img)):
  341. sub_img_id = sorted_subs_info[idx]["sub_img_id"]
  342. sub_img_info_list[sub_img_id]["rec_res"] = rec_res
  343. for sno in range(len(sub_img_info_list)):
  344. rec_res = sub_img_info_list[sno]["rec_res"]
  345. if rec_res["rec_score"] >= text_rec_score_thresh:
  346. single_img_res["rec_texts"].append(rec_res["rec_text"])
  347. single_img_res["rec_scores"].append(rec_res["rec_score"])
  348. single_img_res["rec_polys"].append(dt_polys[sno])
  349. if self.text_type == "general":
  350. rec_boxes = convert_points_to_boxes(single_img_res["rec_polys"])
  351. single_img_res["rec_boxes"] = rec_boxes
  352. else:
  353. single_img_res["rec_boxes"] = np.array([])
  354. yield OCRResult(single_img_res)