__init__.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  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, Union
  16. from importlib import import_module
  17. from ...utils import errors
  18. from ..utils.hpi import HPIConfig
  19. from ..utils.official_models import official_models
  20. from .base import BasePredictor
  21. from .image_classification import ClasPredictor
  22. from .object_detection import DetPredictor
  23. from .keypoint_detection import KptPredictor
  24. from .text_detection import TextDetPredictor
  25. from .text_recognition import TextRecPredictor
  26. from .table_structure_recognition import TablePredictor
  27. from .formula_recognition import FormulaRecPredictor
  28. from .instance_segmentation import InstanceSegPredictor
  29. from .semantic_segmentation import SegPredictor
  30. from .image_feature import ImageFeaturePredictor
  31. from .ts_forecasting import TSFcPredictor
  32. from .ts_anomaly_detection import TSAdPredictor
  33. from .ts_classification import TSClsPredictor
  34. from .image_unwarping import WarpPredictor
  35. from .image_multilabel_classification import MLClasPredictor
  36. from .face_feature import FaceFeaturePredictor
  37. from .open_vocabulary_detection import OVDetPredictor
  38. from .open_vocabulary_segmentation import OVSegPredictor
  39. # from .table_recognition import TablePredictor
  40. # from .general_recognition import ShiTuRecPredictor
  41. from .anomaly_detection import UadPredictor
  42. # from .face_recognition import FaceRecPredictor
  43. from .multilingual_speech_recognition import WhisperPredictor
  44. from .video_classification import VideoClasPredictor
  45. from .video_detection import VideoDetPredictor
  46. module_3d_bev_detection = import_module(".3d_bev_detection", "paddlex.inference.models")
  47. BEVDet3DPredictor = getattr(module_3d_bev_detection, "BEVDet3DPredictor")
  48. def create_predictor(
  49. model_name: str,
  50. model_dir: Optional[str] = None,
  51. device=None,
  52. pp_option=None,
  53. use_hpip: bool = False,
  54. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  55. *args,
  56. **kwargs,
  57. ) -> BasePredictor:
  58. if model_dir is None:
  59. model_dir = check_model(model_name)
  60. else:
  61. assert Path(model_dir).exists(), f"{model_dir} is not exists!"
  62. model_dir = Path(model_dir)
  63. config = BasePredictor.load_config(model_dir)
  64. assert (
  65. model_name == config["Global"]["model_name"]
  66. ), f"Model name mismatch,please input the correct model dir."
  67. return BasePredictor.get(model_name)(
  68. model_dir=model_dir,
  69. config=config,
  70. device=device,
  71. pp_option=pp_option,
  72. use_hpip=use_hpip,
  73. hpi_config=hpi_config,
  74. *args,
  75. **kwargs,
  76. )
  77. def check_model(model):
  78. if Path(model).exists():
  79. return Path(model)
  80. elif model in official_models:
  81. return official_models[model]
  82. else:
  83. raise Exception(
  84. f"The model ({model}) is no exists! Please using directory of local model files or model name supported by PaddleX!"
  85. )