predictor.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  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 os
  15. from copy import deepcopy
  16. from abc import ABC, abstractmethod
  17. from .kernel_option import PaddleInferenceOption
  18. from .utils.paddle_inference_predictor import _PaddleInferencePredictor
  19. from .utils.mixin import FromDictMixin
  20. from .utils.batch import batchable_method, Batcher
  21. from .utils.node import Node
  22. from .utils.official_models import official_models
  23. from ....utils.device import get_device
  24. from ....utils import logging
  25. from ....utils.config import AttrDict
  26. class BasePredictor(ABC, FromDictMixin, Node):
  27. """Base Predictor"""
  28. __is_base = True
  29. MODEL_FILE_TAG = "inference"
  30. def __init__(
  31. self,
  32. model_name,
  33. model_dir,
  34. kernel_option,
  35. output,
  36. pre_transforms=None,
  37. post_transforms=None,
  38. ):
  39. super().__init__()
  40. self.model_name = model_name
  41. self.model_dir = model_dir
  42. self.kernel_option = kernel_option
  43. self.output = output
  44. self.other_src = self.load_other_src()
  45. logging.debug(
  46. f"-------------------- {self.__class__.__name__} --------------------\n\
  47. Model: {self.model_dir}\n\
  48. Env: {self.kernel_option}"
  49. )
  50. self.pre_tfs, self.post_tfs = self.build_transforms(
  51. pre_transforms, post_transforms
  52. )
  53. self._predictor = _PaddleInferencePredictor(
  54. model_dir=model_dir, model_prefix=self.MODEL_FILE_TAG, option=kernel_option
  55. )
  56. def build_transforms(self, pre_transforms, post_transforms):
  57. """build pre-transforms and post-transforms"""
  58. pre_tfs = (
  59. pre_transforms
  60. if pre_transforms is not None
  61. else self._get_pre_transforms_from_config()
  62. )
  63. logging.debug(f"Preprocess Ops: {self._format_transforms(pre_tfs)}")
  64. post_tfs = (
  65. post_transforms
  66. if post_transforms is not None
  67. else self._get_post_transforms_from_config()
  68. )
  69. logging.debug(f"Postprocessing: {self._format_transforms(post_tfs)}")
  70. return pre_tfs, post_tfs
  71. def predict(self, input, batch_size=1):
  72. """predict"""
  73. if not isinstance(input, dict) and not (
  74. isinstance(input, list) and all(isinstance(ele, dict) for ele in input)
  75. ):
  76. raise TypeError(f"`input` should be a dict or a list of dicts.")
  77. orig_input = input
  78. if isinstance(input, dict):
  79. input = [input]
  80. output = []
  81. for mini_batch in Batcher(input, batch_size=batch_size):
  82. mini_batch = self._preprocess(mini_batch, pre_transforms=self.pre_tfs)
  83. for data in mini_batch:
  84. self.check_input_keys(data)
  85. mini_batch = self._run(batch_input=mini_batch)
  86. for data in mini_batch:
  87. self.check_output_keys(data)
  88. mini_batch = self._postprocess(mini_batch, post_transforms=self.post_tfs)
  89. output.extend(mini_batch)
  90. if isinstance(orig_input, dict):
  91. return output[0]
  92. else:
  93. return output
  94. @abstractmethod
  95. def _run(self, batch_input):
  96. raise NotImplementedError
  97. @abstractmethod
  98. def _get_pre_transforms_from_config(self):
  99. """get preprocess transforms"""
  100. raise NotImplementedError
  101. @abstractmethod
  102. def _get_post_transforms_from_config(self):
  103. """get postprocess transforms"""
  104. raise NotImplementedError
  105. @batchable_method
  106. def _preprocess(self, data, pre_transforms):
  107. """preprocess"""
  108. for tf in pre_transforms:
  109. data = tf(data)
  110. return data
  111. @batchable_method
  112. def _postprocess(self, data, post_transforms):
  113. """postprocess"""
  114. for tf in post_transforms:
  115. data = tf(data)
  116. return data
  117. def _format_transforms(self, transforms):
  118. """format transforms"""
  119. ops_str = ", ".join([str(tf) for tf in transforms])
  120. return f"[{ops_str}]"
  121. def load_other_src(self):
  122. """load other source"""
  123. return None
  124. def get_input_keys(self):
  125. """get keys of input dict"""
  126. return self.pre_tfs[0].get_input_keys()
  127. class PredictorBuilderByConfig(object):
  128. """build model predictor"""
  129. def __init__(self, config):
  130. """
  131. Args:
  132. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  133. """
  134. model_name = config.Global.model
  135. device = config.Global.device
  136. predict_config = deepcopy(config.Predict)
  137. model_dir = predict_config.pop("model_dir")
  138. kernel_setting = predict_config.pop("kernel_option", {})
  139. kernel_setting.setdefault("device", device)
  140. kernel_option = PaddleInferenceOption(**kernel_setting)
  141. self.input_path = predict_config.pop("input_path")
  142. self.predictor = BasePredictor.get(model_name)(
  143. model_name=model_name,
  144. model_dir=model_dir,
  145. kernel_option=kernel_option,
  146. output=config.Global.output,
  147. **predict_config,
  148. )
  149. def predict(self):
  150. """predict"""
  151. self.predictor.predict({"input_path": self.input_path})
  152. def build_predictor(*args, **kwargs):
  153. """build predictor by config for dev"""
  154. return PredictorBuilderByConfig(*args, **kwargs)
  155. def create_model(
  156. model_name,
  157. model_dir=None,
  158. kernel_option=None,
  159. output="./",
  160. pre_transforms=None,
  161. post_transforms=None,
  162. *args,
  163. **kwargs,
  164. ):
  165. """create model for predicting using inference model"""
  166. kernel_option = PaddleInferenceOption() if kernel_option is None else kernel_option
  167. if model_dir is None:
  168. if model_name in official_models:
  169. model_dir = official_models[model_name]
  170. return BasePredictor.get(model_name)(
  171. model_name=model_name,
  172. model_dir=model_dir,
  173. kernel_option=kernel_option,
  174. output=output,
  175. pre_transforms=pre_transforms,
  176. post_transforms=post_transforms,
  177. *args,
  178. **kwargs,
  179. )