predictor.py 3.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  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. import numpy as np
  16. from pathlib import Path
  17. from ...base import BasePredictor
  18. from ...base.predictor.transforms import image_common
  19. from .keys import ClsKeys as K
  20. from .utils import InnerConfig
  21. from ....utils import logging
  22. from . import transforms as T
  23. from ..support_models import SUPPORT_MODELS
  24. class ClsPredictor(BasePredictor):
  25. """ Clssification Predictor """
  26. support_models = SUPPORT_MODELS
  27. def load_other_src(self):
  28. """ load the inner config file """
  29. infer_cfg_file_path = os.path.join(self.model_dir, 'inference.yml')
  30. if not os.path.exists(infer_cfg_file_path):
  31. raise FileNotFoundError(
  32. f"Cannot find config file: {infer_cfg_file_path}")
  33. return InnerConfig(infer_cfg_file_path)
  34. @classmethod
  35. def get_input_keys(cls):
  36. """ get input keys """
  37. return [[K.IMAGE], [K.IM_PATH]]
  38. @classmethod
  39. def get_output_keys(cls):
  40. """ get output keys """
  41. return [K.CLS_PRED]
  42. def _run(self, batch_input):
  43. """ run """
  44. input_dict = {}
  45. input_dict[K.IMAGE] = np.stack(
  46. [data[K.IMAGE] for data in batch_input], axis=0).astype(
  47. dtype=np.float32, copy=False)
  48. input_ = [input_dict[K.IMAGE]]
  49. outputs = self._predictor.predict(input_)
  50. cls_outs = outputs[0]
  51. # In-place update
  52. pred = batch_input
  53. for dict_, cls_out in zip(pred, cls_outs):
  54. dict_[K.CLS_PRED] = cls_out
  55. return pred
  56. def _get_pre_transforms_for_data(self, data):
  57. """ get preprocess transforms """
  58. if K.IMAGE not in data:
  59. if K.IM_PATH not in data:
  60. raise KeyError(
  61. f"Key {repr(K.IM_PATH)} is required, but not found.")
  62. logging.info(
  63. f"Transformation operators for data preprocessing will be inferred from config file."
  64. )
  65. pre_transforms = self.other_src.pre_transforms
  66. pre_transforms.insert(0, image_common.ReadImage(format='RGB'))
  67. else:
  68. raise RuntimeError(
  69. f"`{self.__class__.__name__}` does not have default transformation operators to preprocess the input. "
  70. f"Please set `pre_transforms` when using the {repr(K.IMAGE)} key in input dict."
  71. )
  72. pre_transforms.insert(0, T.LoadLabels(self.other_src.labels))
  73. return pre_transforms
  74. def _get_post_transforms_for_data(self, data):
  75. """ get postprocess transforms """
  76. if data.get('cli_flag', False):
  77. output_dir = data.get("output_dir", "./")
  78. return [
  79. # T.SaveDetResults(output_dir, labels=self.other_src.labels),
  80. T.Topk(topk=1),
  81. T.PrintResult()
  82. ]
  83. return []