predictor.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import numpy as np
  13. from pathlib import Path
  14. from ...base import BasePredictor
  15. from ...base.predictor.transforms import image_common
  16. from .keys import ClsKeys as K
  17. from .utils import InnerConfig
  18. from ....utils import logging
  19. from . import transforms as T
  20. from ..model_list import MODELS
  21. class ClsPredictor(BasePredictor):
  22. """ Clssification Predictor """
  23. entities = MODELS
  24. def load_other_src(self):
  25. """ load the inner config file """
  26. infer_cfg_file_path = os.path.join(self.model_dir, 'inference.yml')
  27. if not os.path.exists(infer_cfg_file_path):
  28. raise FileNotFoundError(
  29. f"Cannot find config file: {infer_cfg_file_path}")
  30. return InnerConfig(infer_cfg_file_path)
  31. @classmethod
  32. def get_input_keys(cls):
  33. """ get input keys """
  34. return [[K.IMAGE], [K.IM_PATH]]
  35. @classmethod
  36. def get_output_keys(cls):
  37. """ get output keys """
  38. return [K.CLS_PRED]
  39. def _run(self, batch_input):
  40. """ run """
  41. input_dict = {}
  42. input_dict[K.IMAGE] = np.stack(
  43. [data[K.IMAGE] for data in batch_input], axis=0).astype(
  44. dtype=np.float32, copy=False)
  45. input_ = [input_dict[K.IMAGE]]
  46. outputs = self._predictor.predict(input_)
  47. cls_outs = outputs[0]
  48. # In-place update
  49. pred = batch_input
  50. for dict_, cls_out in zip(pred, cls_outs):
  51. dict_[K.CLS_PRED] = cls_out
  52. return pred
  53. def _get_pre_transforms_from_config(self):
  54. """ get preprocess transforms """
  55. logging.info(
  56. f"Transformation operators for data preprocessing will be inferred from config file."
  57. )
  58. pre_transforms = self.other_src.pre_transforms
  59. pre_transforms.insert(0, image_common.ReadImage(format='RGB'))
  60. return pre_transforms
  61. def _get_post_transforms_from_config(self):
  62. """ get postprocess transforms """
  63. post_transforms = self.other_src.post_transforms
  64. post_transforms.extend([
  65. T.PrintResult(), T.SaveClsResults(self.output,
  66. self.other_src.labels)
  67. ])
  68. return post_transforms