__init__.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  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 pathlib import Path
  15. from typing import Any, Dict, Optional
  16. from ...utils import errors
  17. from ..utils.official_models import official_models
  18. from .base import BasePredictor, BasicPredictor
  19. from .image_classification import ClasPredictor
  20. from .object_detection import DetPredictor
  21. from .keypoint_detection import KptPredictor
  22. from .text_detection import TextDetPredictor
  23. from .text_recognition import TextRecPredictor
  24. from .table_structure_recognition import TablePredictor
  25. from .formula_recognition import FormulaRecPredictor
  26. from .instance_segmentation import InstanceSegPredictor
  27. from .semantic_segmentation import SegPredictor
  28. from .image_feature import ImageFeaturePredictor
  29. from .ts_forecasting import TSFcPredictor
  30. from .ts_anomaly_detection import TSAdPredictor
  31. from .ts_classification import TSClsPredictor
  32. from .image_unwarping import WarpPredictor
  33. from .image_multilabel_classification import MLClasPredictor
  34. # from .table_recognition import TablePredictor
  35. # from .general_recognition import ShiTuRecPredictor
  36. from .anomaly_detection import UadPredictor
  37. # from .face_recognition import FaceRecPredictor
  38. from .multilingual_speech_recognition import WhisperPredictor
  39. from .video_classification import VideoClasPredictor
  40. def _create_hp_predictor(
  41. model_name, model_dir, device, config, hpi_params, *args, **kwargs
  42. ):
  43. try:
  44. from paddlex_hpi.models import HPPredictor
  45. except ModuleNotFoundError:
  46. raise RuntimeError(
  47. "The PaddleX HPI plugin is not properly installed, and the high-performance model inference features are not available."
  48. ) from None
  49. try:
  50. predictor = HPPredictor.get(model_name)(
  51. model_dir=model_dir,
  52. config=config,
  53. device=device,
  54. *args,
  55. hpi_params=hpi_params,
  56. **kwargs,
  57. )
  58. except errors.others.ClassNotFoundException:
  59. raise ValueError(
  60. f"{model_name} is not supported by the PaddleX HPI plugin."
  61. ) from None
  62. return predictor
  63. def create_predictor(
  64. model: str,
  65. device=None,
  66. pp_option=None,
  67. use_hpip: bool = False,
  68. hpi_params: Optional[Dict[str, Any]] = None,
  69. *args,
  70. **kwargs,
  71. ) -> BasePredictor:
  72. model_dir = check_model(model)
  73. config = BasePredictor.load_config(model_dir)
  74. model_name = config["Global"]["model_name"]
  75. if use_hpip:
  76. return _create_hp_predictor(
  77. model_name=model_name,
  78. model_dir=model_dir,
  79. config=config,
  80. hpi_params=hpi_params,
  81. device=device,
  82. *args,
  83. **kwargs,
  84. )
  85. else:
  86. return BasicPredictor.get(model_name)(
  87. model_dir=model_dir,
  88. config=config,
  89. device=device,
  90. pp_option=pp_option,
  91. *args,
  92. **kwargs,
  93. )
  94. def check_model(model):
  95. if Path(model).exists():
  96. return Path(model)
  97. elif model in official_models:
  98. return official_models[model]
  99. else:
  100. raise Exception(
  101. f"The model ({model}) is no exists! Please using directory of local model files or model name supported by PaddleX!"
  102. )