pipeline.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  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 rotate_image
  23. from .result import DocPreprocessorResult
  24. class DocPreprocessorPipeline(BasePipeline):
  25. """Doc Preprocessor Pipeline"""
  26. entities = "doc_preprocessor"
  27. def __init__(
  28. self,
  29. config: Dict,
  30. device: Optional[str] = None,
  31. pp_option: Optional[PaddlePredictorOption] = None,
  32. use_hpip: bool = False,
  33. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  34. ) -> None:
  35. """Initializes the doc preprocessor pipeline.
  36. Args:
  37. config (Dict): Configuration dictionary containing various settings.
  38. device (str, optional): Device to run the predictions on. Defaults to None.
  39. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  40. use_hpip (bool, optional): Whether to use the high-performance
  41. inference plugin (HPIP) by default. Defaults to False.
  42. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  43. The default high-performance inference configuration dictionary.
  44. Defaults to None.
  45. """
  46. super().__init__(
  47. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  48. )
  49. self.use_doc_orientation_classify = config.get(
  50. "use_doc_orientation_classify", True
  51. )
  52. if self.use_doc_orientation_classify:
  53. doc_ori_classify_config = config.get("SubModules", {}).get(
  54. "DocOrientationClassify",
  55. {"model_config_error": "config error for doc_ori_classify_model!"},
  56. )
  57. self.doc_ori_classify_model = self.create_model(doc_ori_classify_config)
  58. self.use_doc_unwarping = config.get("use_doc_unwarping", True)
  59. if self.use_doc_unwarping:
  60. doc_unwarping_config = config.get("SubModules", {}).get(
  61. "DocUnwarping",
  62. {"model_config_error": "config error for doc_unwarping_model!"},
  63. )
  64. self.doc_unwarping_model = self.create_model(doc_unwarping_config)
  65. self.batch_sampler = ImageBatchSampler(batch_size=1)
  66. self.img_reader = ReadImage(format="BGR")
  67. def check_model_settings_valid(self, model_settings: Dict) -> bool:
  68. """
  69. Check if the the input params for model settings are valid based on the initialized models.
  70. Args:
  71. model_settings (Dict): A dictionary containing model settings.
  72. Returns:
  73. bool: True if all required models are initialized according to the model settings, False otherwise.
  74. """
  75. if (
  76. model_settings["use_doc_orientation_classify"]
  77. and not self.use_doc_orientation_classify
  78. ):
  79. logging.error(
  80. "Set use_doc_orientation_classify, but the model for doc orientation classify is not initialized."
  81. )
  82. return False
  83. if model_settings["use_doc_unwarping"] and not self.use_doc_unwarping:
  84. logging.error(
  85. "Set use_doc_unwarping, but the model for doc unwarping is not initialized."
  86. )
  87. return False
  88. return True
  89. def get_model_settings(
  90. self, use_doc_orientation_classify, use_doc_unwarping
  91. ) -> dict:
  92. """
  93. Retrieve the model settings dictionary based on input parameters.
  94. Args:
  95. use_doc_orientation_classify (bool, optional): Whether to use document orientation classification.
  96. use_doc_unwarping (bool, optional): Whether to use document unwarping.
  97. Returns:
  98. dict: A dictionary containing the model settings.
  99. """
  100. if use_doc_orientation_classify is None:
  101. use_doc_orientation_classify = self.use_doc_orientation_classify
  102. if use_doc_unwarping is None:
  103. use_doc_unwarping = self.use_doc_unwarping
  104. model_settings = {
  105. "use_doc_orientation_classify": use_doc_orientation_classify,
  106. "use_doc_unwarping": use_doc_unwarping,
  107. }
  108. return model_settings
  109. def predict(
  110. self,
  111. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  112. use_doc_orientation_classify: Optional[bool] = None,
  113. use_doc_unwarping: Optional[bool] = None,
  114. ) -> DocPreprocessorResult:
  115. """
  116. Predict the preprocessing result for the input image or images.
  117. Args:
  118. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images or pdfs.
  119. use_doc_orientation_classify (bool): Whether to use document orientation classification.
  120. use_doc_unwarping (bool): Whether to use document unwarping.
  121. **kwargs: Additional keyword arguments.
  122. Returns:
  123. DocPreprocessorResult: A generator yielding preprocessing results.
  124. """
  125. model_settings = self.get_model_settings(
  126. use_doc_orientation_classify, use_doc_unwarping
  127. )
  128. if not self.check_model_settings_valid(model_settings):
  129. yield {"error": "the input params for model settings are invalid!"}
  130. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  131. image_array = self.img_reader(batch_data.instances)[0]
  132. if model_settings["use_doc_orientation_classify"]:
  133. pred = next(self.doc_ori_classify_model(image_array))
  134. angle = int(pred["label_names"][0])
  135. rot_img = rotate_image(image_array, angle)
  136. else:
  137. angle = -1
  138. rot_img = image_array
  139. if model_settings["use_doc_unwarping"]:
  140. output_img = next(self.doc_unwarping_model(rot_img))["doctr_img"]
  141. else:
  142. output_img = rot_img
  143. single_img_res = {
  144. "input_path": batch_data.input_paths[0],
  145. "page_index": batch_data.page_indexes[0],
  146. "input_img": image_array,
  147. "model_settings": model_settings,
  148. "angle": angle,
  149. "rot_img": rot_img,
  150. "output_img": output_img,
  151. }
  152. yield DocPreprocessorResult(single_img_res)