predictor.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  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 ....utils import logging
  14. from ...base import BasePredictor
  15. from ...base.predictor.transforms import image_common
  16. from . import transforms as T
  17. from .keys import DetKeys as K
  18. from .utils import InnerConfig
  19. from ..model_list import MODELS
  20. class DetPredictor(BasePredictor):
  21. """ Detection Predictor """
  22. entities = MODELS
  23. def load_other_src(self):
  24. """ load the inner config file """
  25. infer_cfg_file_path = os.path.join(self.model_dir, 'inference.yml')
  26. if not os.path.exists(infer_cfg_file_path):
  27. raise FileNotFoundError(
  28. f"Cannot find config file: {infer_cfg_file_path}")
  29. return InnerConfig(infer_cfg_file_path)
  30. @classmethod
  31. def get_input_keys(cls):
  32. """ get input keys """
  33. return [[K.IMAGE], [K.IM_PATH]]
  34. @classmethod
  35. def get_output_keys(cls):
  36. """ get output keys """
  37. return [K.BOXES]
  38. def _run(self, batch_input):
  39. """ run """
  40. input_dict = {}
  41. input_dict["image"] = np.stack(
  42. [data[K.IMAGE] for data in batch_input], axis=0).astype(
  43. dtype=np.float32, copy=False)
  44. input_dict["scale_factor"] = np.stack(
  45. [data[K.SCALE_FACTOR][::-1] for data in batch_input],
  46. axis=0).astype(
  47. dtype=np.float32, copy=False)
  48. input_dict["im_shape"] = np.stack(
  49. [data[K.IM_SIZE][::-1] for data in batch_input], axis=0).astype(
  50. dtype=np.float32, copy=False)
  51. input_ = [input_dict[i] for i in self._predictor.get_input_names()]
  52. batch_np_boxes, batch_np_boxes_num = self._predictor.predict(input_)
  53. pred = batch_input
  54. box_idx_start = 0
  55. for idx in range(len(batch_input)):
  56. np_boxes_num = batch_np_boxes_num[idx]
  57. box_idx_end = box_idx_start + np_boxes_num
  58. np_boxes = batch_np_boxes[box_idx_start:box_idx_end]
  59. box_idx_start = box_idx_end
  60. batch_input[idx][K.BOXES] = np_boxes
  61. return pred
  62. def _get_pre_transforms_from_config(self):
  63. """ get preprocess transforms """
  64. logging.info(
  65. f"Transformation operators for data preprocessing will be inferred from config file."
  66. )
  67. pre_transforms = self.other_src.pre_transforms
  68. pre_transforms.insert(0, image_common.ReadImage(format='RGB'))
  69. return pre_transforms
  70. def _get_post_transforms_from_config(self):
  71. """ get postprocess transforms """
  72. return [
  73. T.SaveDetResults(
  74. save_dir=self.output, labels=self.other_src.labels),
  75. T.PrintResult()
  76. ]