cli.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. import os
  2. from pathlib import Path
  3. import click
  4. from loguru import logger
  5. import magic_pdf.model as model_config
  6. from magic_pdf.libs.version import __version__
  7. from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
  8. from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
  9. from magic_pdf.tools.common import do_parse, parse_pdf_methods
  10. @click.command()
  11. @click.version_option(__version__,
  12. '--version',
  13. '-v',
  14. help='display the version and exit')
  15. @click.option(
  16. '-p',
  17. '--path',
  18. 'path',
  19. type=click.Path(exists=True),
  20. required=True,
  21. help='local pdf filepath or directory',
  22. )
  23. @click.option(
  24. '-o',
  25. '--output-dir',
  26. 'output_dir',
  27. type=click.Path(),
  28. required=True,
  29. help='output local directory',
  30. default='',
  31. )
  32. @click.option(
  33. '-m',
  34. '--method',
  35. 'method',
  36. type=parse_pdf_methods,
  37. help="""the method for parsing pdf.
  38. ocr: using ocr technique to extract information from pdf.
  39. txt: suitable for the text-based pdf only and outperform ocr.
  40. auto: automatically choose the best method for parsing pdf from ocr and txt.
  41. without method specified, auto will be used by default.""",
  42. default='auto',
  43. )
  44. def cli(path, output_dir, method):
  45. model_config.__use_inside_model__ = True
  46. model_config.__model_mode__ = 'full'
  47. os.makedirs(output_dir, exist_ok=True)
  48. def read_fn(path):
  49. disk_rw = DiskReaderWriter(os.path.dirname(path))
  50. return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
  51. def parse_doc(doc_path: str):
  52. try:
  53. file_name = str(Path(doc_path).stem)
  54. pdf_data = read_fn(doc_path)
  55. do_parse(
  56. output_dir,
  57. file_name,
  58. pdf_data,
  59. [],
  60. method,
  61. )
  62. except Exception as e:
  63. logger.exception(e)
  64. if os.path.isdir(path):
  65. for doc_path in Path(path).glob('*.pdf'):
  66. parse_doc(doc_path)
  67. else:
  68. parse_doc(path)
  69. if __name__ == '__main__':
  70. cli()