pipeline.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277
  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.object_detection.result import DetResult
  20. from ...utils.hpi import HPIConfig
  21. from ...utils.pp_option import PaddlePredictorOption
  22. from ..base import BasePipeline
  23. from ..components import CropByBoxes
  24. from .result import SealRecognitionResult
  25. class SealRecognitionPipeline(BasePipeline):
  26. """Seal Recognition Pipeline"""
  27. entities = ["seal_recognition"]
  28. def __init__(
  29. self,
  30. config: Dict,
  31. device: str = None,
  32. pp_option: PaddlePredictorOption = None,
  33. use_hpip: bool = False,
  34. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  35. ) -> None:
  36. """Initializes the seal recognition pipeline.
  37. Args:
  38. config (Dict): Configuration dictionary containing various settings.
  39. device (str, optional): Device to run the predictions on. Defaults to None.
  40. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  41. use_hpip (bool, optional): Whether to use the high-performance
  42. inference plugin (HPIP) by default. Defaults to False.
  43. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  44. The default high-performance inference configuration dictionary.
  45. Defaults to None.
  46. """
  47. super().__init__(
  48. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  49. )
  50. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  51. if self.use_doc_preprocessor:
  52. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  53. "DocPreprocessor",
  54. {
  55. "pipeline_config_error": "config error for doc_preprocessor_pipeline!"
  56. },
  57. )
  58. self.doc_preprocessor_pipeline = self.create_pipeline(
  59. doc_preprocessor_config
  60. )
  61. self.use_layout_detection = config.get("use_layout_detection", True)
  62. if self.use_layout_detection:
  63. layout_det_config = config.get("SubModules", {}).get(
  64. "LayoutDetection",
  65. {"model_config_error": "config error for layout_det_model!"},
  66. )
  67. layout_kwargs = {}
  68. if (threshold := layout_det_config.get("threshold", None)) is not None:
  69. layout_kwargs["threshold"] = threshold
  70. if (layout_nms := layout_det_config.get("layout_nms", None)) is not None:
  71. layout_kwargs["layout_nms"] = layout_nms
  72. if (
  73. layout_unclip_ratio := layout_det_config.get(
  74. "layout_unclip_ratio", None
  75. )
  76. ) is not None:
  77. layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio
  78. if (
  79. layout_merge_bboxes_mode := layout_det_config.get(
  80. "layout_merge_bboxes_mode", None
  81. )
  82. ) is not None:
  83. layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode
  84. self.layout_det_model = self.create_model(
  85. layout_det_config, **layout_kwargs
  86. )
  87. seal_ocr_config = config.get("SubPipelines", {}).get(
  88. "SealOCR", {"pipeline_config_error": "config error for seal_ocr_pipeline!"}
  89. )
  90. self.seal_ocr_pipeline = self.create_pipeline(seal_ocr_config)
  91. self._crop_by_boxes = CropByBoxes()
  92. self.batch_sampler = ImageBatchSampler(batch_size=1)
  93. self.img_reader = ReadImage(format="BGR")
  94. def check_model_settings_valid(
  95. self, model_settings: Dict, layout_det_res: DetResult
  96. ) -> bool:
  97. """
  98. Check if the input parameters are valid based on the initialized models.
  99. Args:
  100. model_settings (Dict): A dictionary containing input parameters.
  101. layout_det_res (DetResult): Layout detection result.
  102. Returns:
  103. bool: True if all required models are initialized according to input parameters, False otherwise.
  104. """
  105. if model_settings["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  106. logging.error(
  107. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  108. )
  109. return False
  110. if model_settings["use_layout_detection"]:
  111. if layout_det_res is not None:
  112. logging.error(
  113. "The layout detection model has already been initialized, please set use_layout_detection=False"
  114. )
  115. return False
  116. if not self.use_layout_detection:
  117. logging.error(
  118. "Set use_layout_detection, but the models for layout detection are not initialized."
  119. )
  120. return False
  121. return True
  122. def get_model_settings(
  123. self,
  124. use_doc_orientation_classify: Optional[bool],
  125. use_doc_unwarping: Optional[bool],
  126. use_layout_detection: Optional[bool],
  127. ) -> dict:
  128. """
  129. Get the model settings based on the provided parameters or default values.
  130. Args:
  131. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  132. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  133. use_layout_detection (Optional[bool]): Whether to use layout detection.
  134. Returns:
  135. dict: A dictionary containing the model settings.
  136. """
  137. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  138. use_doc_preprocessor = self.use_doc_preprocessor
  139. else:
  140. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  141. use_doc_preprocessor = True
  142. else:
  143. use_doc_preprocessor = False
  144. if use_layout_detection is None:
  145. use_layout_detection = self.use_layout_detection
  146. return dict(
  147. use_doc_preprocessor=use_doc_preprocessor,
  148. use_layout_detection=use_layout_detection,
  149. )
  150. def predict(
  151. self,
  152. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  153. use_doc_orientation_classify: Optional[bool] = None,
  154. use_doc_unwarping: Optional[bool] = None,
  155. use_layout_detection: Optional[bool] = None,
  156. layout_det_res: Optional[DetResult] = None,
  157. layout_threshold: Optional[Union[float, dict]] = None,
  158. layout_nms: Optional[bool] = None,
  159. layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
  160. layout_merge_bboxes_mode: Optional[str] = None,
  161. seal_det_limit_side_len: Optional[int] = None,
  162. seal_det_limit_type: Optional[str] = None,
  163. seal_det_thresh: Optional[float] = None,
  164. seal_det_box_thresh: Optional[float] = None,
  165. seal_det_unclip_ratio: Optional[float] = None,
  166. seal_rec_score_thresh: Optional[float] = None,
  167. **kwargs,
  168. ) -> SealRecognitionResult:
  169. model_settings = self.get_model_settings(
  170. use_doc_orientation_classify, use_doc_unwarping, use_layout_detection
  171. )
  172. if not self.check_model_settings_valid(model_settings, layout_det_res):
  173. yield {"error": "the input params for model settings are invalid!"}
  174. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  175. image_array = self.img_reader(batch_data.instances)[0]
  176. if model_settings["use_doc_preprocessor"]:
  177. doc_preprocessor_res = next(
  178. self.doc_preprocessor_pipeline(
  179. image_array,
  180. use_doc_orientation_classify=use_doc_orientation_classify,
  181. use_doc_unwarping=use_doc_unwarping,
  182. )
  183. )
  184. else:
  185. doc_preprocessor_res = {"output_img": image_array}
  186. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  187. seal_res_list = []
  188. seal_region_id = 1
  189. if not model_settings["use_layout_detection"] and layout_det_res is None:
  190. layout_det_res = {}
  191. seal_ocr_res = next(
  192. self.seal_ocr_pipeline(
  193. doc_preprocessor_image,
  194. text_det_limit_side_len=seal_det_limit_side_len,
  195. text_det_limit_type=seal_det_limit_type,
  196. text_det_thresh=seal_det_thresh,
  197. text_det_box_thresh=seal_det_box_thresh,
  198. text_det_unclip_ratio=seal_det_unclip_ratio,
  199. text_rec_score_thresh=seal_rec_score_thresh,
  200. )
  201. )
  202. seal_ocr_res["seal_region_id"] = seal_region_id
  203. seal_res_list.append(seal_ocr_res)
  204. seal_region_id += 1
  205. else:
  206. if model_settings["use_layout_detection"]:
  207. layout_det_res = next(
  208. self.layout_det_model(
  209. doc_preprocessor_image,
  210. threshold=layout_threshold,
  211. layout_nms=layout_nms,
  212. layout_unclip_ratio=layout_unclip_ratio,
  213. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  214. )
  215. )
  216. for box_info in layout_det_res["boxes"]:
  217. if box_info["label"].lower() in ["seal"]:
  218. crop_img_info = self._crop_by_boxes(
  219. doc_preprocessor_image, [box_info]
  220. )
  221. crop_img_info = crop_img_info[0]
  222. seal_ocr_res = next(
  223. self.seal_ocr_pipeline(
  224. crop_img_info["img"],
  225. text_det_limit_side_len=seal_det_limit_side_len,
  226. text_det_limit_type=seal_det_limit_type,
  227. text_det_thresh=seal_det_thresh,
  228. text_det_box_thresh=seal_det_box_thresh,
  229. text_det_unclip_ratio=seal_det_unclip_ratio,
  230. text_rec_score_thresh=seal_rec_score_thresh,
  231. )
  232. )
  233. seal_ocr_res["seal_region_id"] = seal_region_id
  234. seal_res_list.append(seal_ocr_res)
  235. seal_region_id += 1
  236. single_img_res = {
  237. "input_path": batch_data.input_paths[0],
  238. "page_index": batch_data.page_indexes[0],
  239. "doc_preprocessor_res": doc_preprocessor_res,
  240. "layout_det_res": layout_det_res,
  241. "seal_res_list": seal_res_list,
  242. "model_settings": model_settings,
  243. }
  244. yield SealRecognitionResult(single_img_res)