command.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. from six import text_type as _text_type
  2. import argparse
  3. import sys
  4. def arg_parser():
  5. parser = argparse.ArgumentParser()
  6. parser.add_argument(
  7. "--model_dir",
  8. "-m",
  9. type=_text_type,
  10. default=None,
  11. help="define model directory path")
  12. parser.add_argument(
  13. "--save_dir",
  14. "-s",
  15. type=_text_type,
  16. default=None,
  17. help="path to save inference model")
  18. parser.add_argument(
  19. "--version",
  20. "-v",
  21. action="store_true",
  22. default=False,
  23. help="get version of PaddleX")
  24. parser.add_argument(
  25. "--export_inference",
  26. "-e",
  27. action="store_true",
  28. default=False,
  29. help="export inference model for C++/Python deployment")
  30. return parser
  31. def main():
  32. import os
  33. os.environ['CUDA_VISIBLE_DEVICES'] = ""
  34. import paddlex as pdx
  35. if len(sys.argv) < 2:
  36. print("Use command 'paddlex -h` to print the help information\n")
  37. return
  38. parser = arg_parser()
  39. args = parser.parse_args()
  40. if args.version:
  41. print("PaddleX-{}".format(pdx.__version__))
  42. print("Repo: https://github.com/PaddlePaddle/PaddleX.git")
  43. print("Email: paddlex@baidu.com")
  44. return
  45. if args.export_inference:
  46. assert args.model_dir is not None, "--model_dir should be defined while exporting inference model"
  47. assert args.save_dir is not None, "--save_dir should be defined to save inference model"
  48. model = pdx.load_model(args.model_dir)
  49. model.export_inference_model(args.save_dir)
  50. if __name__ == "__main__":
  51. main()