cli.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  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. @click.option(
  45. "-d",
  46. "--debug",
  47. "debug_able",
  48. type=bool,
  49. help="Enables detailed debugging information during the execution of the CLI commands.",
  50. default=False,
  51. )
  52. def cli(path, output_dir, method, debug_able):
  53. model_config.__use_inside_model__ = True
  54. model_config.__model_mode__ = 'full'
  55. os.makedirs(output_dir, exist_ok=True)
  56. def read_fn(path):
  57. disk_rw = DiskReaderWriter(os.path.dirname(path))
  58. return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
  59. def parse_doc(doc_path: str):
  60. try:
  61. file_name = str(Path(doc_path).stem)
  62. pdf_data = read_fn(doc_path)
  63. do_parse(
  64. output_dir,
  65. file_name,
  66. pdf_data,
  67. [],
  68. method,
  69. debug_able,
  70. )
  71. except Exception as e:
  72. logger.exception(e)
  73. if os.path.isdir(path):
  74. for doc_path in Path(path).glob('*.pdf'):
  75. parse_doc(doc_path)
  76. else:
  77. parse_doc(path)
  78. if __name__ == '__main__':
  79. cli()