model.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  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. from ...base.utils.arg import CLIArgument
  13. from ...base.utils.subprocess import CompletedProcess
  14. from ....utils.misc import abspath
  15. from ..text_rec.model import TextRecModel
  16. class TextDetModel(TextRecModel):
  17. """ Text Detection Model """
  18. def infer(self,
  19. model_dir: str,
  20. input_path: str,
  21. device: str='gpu',
  22. save_dir: str=None,
  23. **kwargs) -> CompletedProcess:
  24. """predict image using infernece model
  25. Args:
  26. model_dir (str): the directory path of inference model files that would use to predict.
  27. input_path (str): the path of image that would be predict.
  28. device (str, optional): the running device. Defaults to 'gpu'.
  29. save_dir (str, optional): the directory path to save output. Defaults to None.
  30. Returns:
  31. CompletedProcess: the result of infering subprocess execution.
  32. """
  33. config = self.config.copy()
  34. cli_args = []
  35. model_dir = abspath(model_dir)
  36. cli_args.append(CLIArgument('--det_model_dir', model_dir))
  37. input_path = abspath(input_path)
  38. cli_args.append(CLIArgument('--image_dir', input_path))
  39. device_type, _ = self.runner.parse_device(device)
  40. cli_args.append(CLIArgument('--use_gpu', str(device_type == 'gpu')))
  41. if save_dir is not None:
  42. save_dir = abspath(save_dir)
  43. else:
  44. # `save_dir` is None
  45. save_dir = abspath(os.path.join('output', 'infer'))
  46. cli_args.append(CLIArgument('--draw_img_save_dir', save_dir))
  47. model_type = config._get_model_type()
  48. cli_args.append(CLIArgument('--det_algorithm', model_type))
  49. self._assert_empty_kwargs(kwargs)
  50. with self._create_new_config_file() as config_path:
  51. config.dump(config_path)
  52. return self.runner.infer(config_path, cli_args, device)