__init__.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  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 .base import BasePredictor
  16. from .image_classification import ClasPredictor
  17. from .text_detection import TextDetPredictor
  18. from .text_recognition import TextRecPredictor
  19. from .table_recognition import TablePredictor
  20. from .official_models import official_models
  21. def create_predictor(
  22. model: str, device: str = None, pp_option=None, *args, **kwargs
  23. ) -> BasePredictor:
  24. model_dir = check_model(model)
  25. config = BasePredictor.load_config(model_dir)
  26. model_name = config["Global"]["model_name"]
  27. return BasePredictor.get(model_name)(
  28. model_dir=model_dir,
  29. config=config,
  30. device=device,
  31. pp_option=pp_option,
  32. *args,
  33. **kwargs,
  34. )
  35. def check_model(model):
  36. if Path(model).exists():
  37. return Path(model)
  38. elif model in official_models:
  39. return official_models[model]
  40. else:
  41. raise Exception(
  42. f"The model ({model}) is no exists! Please using directory of local model files or model name supported by PaddleX!"
  43. )