predictor.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  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. from typing import Any, Union, Dict, List, Tuple
  15. import numpy as np
  16. from ....utils.func_register import FuncRegister
  17. from ....modules.image_classification.model_list import MODELS
  18. from ...common.batch_sampler import ImageBatchSampler
  19. from ...common.reader import ReadImage
  20. from ..common import (
  21. Resize,
  22. ResizeByShort,
  23. Normalize,
  24. ToCHWImage,
  25. ToBatch,
  26. StaticInfer,
  27. )
  28. from ..base import BasicPredictor
  29. from .processors import Crop, Topk
  30. from .result import TopkResult
  31. class ClasPredictor(BasicPredictor):
  32. """ClasPredictor that inherits from BasicPredictor."""
  33. entities = MODELS
  34. _FUNC_MAP = {}
  35. register = FuncRegister(_FUNC_MAP)
  36. def __init__(self, *args: List, **kwargs: Dict) -> None:
  37. """Initializes ClasPredictor.
  38. Args:
  39. *args: Arbitrary positional arguments passed to the superclass.
  40. **kwargs: Arbitrary keyword arguments passed to the superclass.
  41. """
  42. super().__init__(*args, **kwargs)
  43. self.preprocessors, self.infer, self.postprocessors = self._build()
  44. def _build_batch_sampler(self) -> ImageBatchSampler:
  45. """Builds and returns an ImageBatchSampler instance.
  46. Returns:
  47. ImageBatchSampler: An instance of ImageBatchSampler.
  48. """
  49. return ImageBatchSampler()
  50. def _get_result_class(self) -> type:
  51. """Returns the result class, TopkResult.
  52. Returns:
  53. type: The TopkResult class.
  54. """
  55. return TopkResult
  56. def _build(self) -> Tuple:
  57. """Build the preprocessors, inference engine, and postprocessors based on the configuration.
  58. Returns:
  59. tuple: A tuple containing the preprocessors, inference engine, and postprocessors.
  60. """
  61. preprocessors = {"Read": ReadImage(format="RGB")}
  62. for cfg in self.config["PreProcess"]["transform_ops"]:
  63. tf_key = list(cfg.keys())[0]
  64. func = self._FUNC_MAP[tf_key]
  65. args = cfg.get(tf_key, {})
  66. name, op = func(self, **args) if args else func(self)
  67. preprocessors[name] = op
  68. preprocessors["ToBatch"] = ToBatch()
  69. infer = StaticInfer(
  70. model_dir=self.model_dir,
  71. model_prefix=self.MODEL_FILE_PREFIX,
  72. option=self.pp_option,
  73. )
  74. postprocessors = {}
  75. for key in self.config["PostProcess"]:
  76. func = self._FUNC_MAP.get(key)
  77. args = self.config["PostProcess"].get(key, {})
  78. name, op = func(self, **args) if args else func(self)
  79. postprocessors[name] = op
  80. return preprocessors, infer, postprocessors
  81. def process(self, batch_data: List[Union[str, np.ndarray]]) -> Dict[str, Any]:
  82. """
  83. Process a batch of data through the preprocessing, inference, and postprocessing.
  84. Args:
  85. batch_data (List[Union[str, np.ndarray], ...]): A batch of input data (e.g., image file paths).
  86. Returns:
  87. dict: A dictionary containing the input path, raw image, class IDs, scores, and label names for every instance of the batch. Keys include 'input_path', 'input_img', 'class_ids', 'scores', and 'label_names'.
  88. """
  89. batch_raw_imgs = self.preprocessors["Read"](imgs=batch_data)
  90. batch_imgs = self.preprocessors["Resize"](imgs=batch_raw_imgs)
  91. batch_imgs = self.preprocessors["Crop"](imgs=batch_imgs)
  92. batch_imgs = self.preprocessors["Normalize"](imgs=batch_imgs)
  93. batch_imgs = self.preprocessors["ToCHW"](imgs=batch_imgs)
  94. x = self.preprocessors["ToBatch"](imgs=batch_imgs)
  95. batch_preds = self.infer(x=x)
  96. batch_class_ids, batch_scores, batch_label_names = self.postprocessors["Topk"](
  97. batch_preds
  98. )
  99. return {
  100. "input_path": batch_data,
  101. "input_img": batch_raw_imgs,
  102. "class_ids": batch_class_ids,
  103. "scores": batch_scores,
  104. "label_names": batch_label_names,
  105. }
  106. @register("ResizeImage")
  107. # TODO(gaotingquan): backend & interpolation
  108. def build_resize(
  109. self, resize_short=None, size=None, backend="cv2", interpolation="LINEAR"
  110. ):
  111. assert resize_short or size
  112. if resize_short:
  113. op = ResizeByShort(
  114. target_short_edge=resize_short, size_divisor=None, interp="LINEAR"
  115. )
  116. else:
  117. op = Resize(target_size=size)
  118. return "Resize", op
  119. @register("CropImage")
  120. def build_crop(self, size=224):
  121. return "Crop", Crop(crop_size=size)
  122. @register("NormalizeImage")
  123. def build_normalize(
  124. self,
  125. mean=[0.485, 0.456, 0.406],
  126. std=[0.229, 0.224, 0.225],
  127. scale=1 / 255,
  128. order="",
  129. channel_num=3,
  130. ):
  131. assert channel_num == 3
  132. assert order == ""
  133. return "Normalize", Normalize(scale=scale, mean=mean, std=std)
  134. @register("ToCHWImage")
  135. def build_to_chw(self):
  136. return "ToCHW", ToCHWImage()
  137. @register("Topk")
  138. def build_topk(self, topk, label_list=None):
  139. return "Topk", Topk(topk=int(topk), class_ids=label_list)