predictor.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  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. from pathlib import Path
  12. from typing import Union
  13. from ...utils import logging
  14. from ..base.build_model import build_model
  15. from ..base.predictor import BasePredictor
  16. from ...utils.errors import raise_unsupported_api_error, raise_model_not_found_error
  17. from .model_list import MODELS
  18. class TSFCPredictor(BasePredictor):
  19. """ TS Forecast Model Predictor """
  20. entities = MODELS
  21. def __init__(self, config):
  22. """Initialize the instance.
  23. Args:
  24. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  25. """
  26. self.global_config = config.Global
  27. self.predict_config = config.Predict
  28. config_path = self.get_config_path()
  29. self.pdx_config, self.pdx_model = build_model(
  30. self.global_config.model, config_path=config_path)
  31. def get_config_path(self) -> Union[str, None]:
  32. """
  33. get config path
  34. Returns:
  35. config_path (str): The path to the config
  36. """
  37. model_dir = self.predict_config.model_dir
  38. if Path(model_dir).exists():
  39. config_path = Path(model_dir).parent.parent / "config.yaml"
  40. if config_path.exists():
  41. return config_path
  42. else:
  43. logging.warning(
  44. f"The config file(`{config_path}`) related to model weight file(`{self.predict_config.model_dir}`) \
  45. is not exist, use default instead.")
  46. else:
  47. raise_model_not_found_error(model_dir)
  48. return None
  49. def predict(self, input=None, batch_size=1):
  50. """execute model predict
  51. Returns:
  52. dict: the prediction results
  53. """
  54. results = self.predict()
  55. def predict(self):
  56. """predict using specified model
  57. """
  58. # self.update_config()
  59. result = self.pdx_model.predict(**self.get_predict_kwargs())
  60. assert result.returncode == 0, f"Encountered an unexpected error({result.returncode}) in predicting!"
  61. def get_predict_kwargs(self) -> dict:
  62. """get key-value arguments of model predict function
  63. Returns:
  64. dict: the arguments of predict function.
  65. """
  66. return {
  67. "weight_path": self.predict_config.model_dir,
  68. "input_path": self.predict_config.input_path,
  69. "device": self.global_config.device,
  70. "save_dir": self.global_config.output
  71. }
  72. def _get_post_transforms_from_config(self):
  73. pass
  74. def _get_pre_transforms_from_config(self):
  75. pass
  76. def _run(self):
  77. pass
  78. def get_input_keys(self):
  79. """ get input keys """
  80. pass
  81. def get_output_keys(self):
  82. """ get output keys """
  83. pass