Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader | |
| from magic_pdf.data.dataset import PymuDocDataset | |
| from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze | |
| from magic_pdf.config.enums import SupportedPdfParseMethod | |
| def convert_pdf(pdf_file_path): | |
| # Get filename and prepare output paths | |
| pdf_file_name = os.path.basename(pdf_file_path) | |
| name_without_suff = os.path.splitext(pdf_file_name)[0] | |
| # prepare env | |
| local_image_dir, local_md_dir = "output/images", "output" | |
| image_dir = str(os.path.basename(local_image_dir)) | |
| os.makedirs(local_image_dir, exist_ok=True) | |
| os.makedirs(local_md_dir, exist_ok=True) | |
| image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( | |
| local_md_dir | |
| ) | |
| # read bytes | |
| reader1 = FileBasedDataReader(os.path.dirname(pdf_file_path)) | |
| pdf_bytes = reader1.read(pdf_file_name) # read the pdf content | |
| print(f"Processing PDF: {pdf_file_path}") | |
| # proc | |
| ## Create Dataset Instance | |
| ds = PymuDocDataset(pdf_bytes) | |
| ## inference | |
| if ds.classify() == SupportedPdfParseMethod.OCR: | |
| infer_result = ds.apply(doc_analyze, ocr=True) | |
| ## pipeline | |
| pipe_result = infer_result.pipe_ocr_mode(image_writer) | |
| else: | |
| infer_result = ds.apply(doc_analyze, ocr=False) | |
| ## pipeline | |
| pipe_result = infer_result.pipe_txt_mode(image_writer) | |
| ### draw model result on each page | |
| infer_result.draw_model(os.path.join(local_md_dir, f"{name_without_suff}_model.pdf")) | |
| ### get model inference result | |
| model_inference_result = infer_result.get_infer_res() | |
| ### draw layout result on each page | |
| pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf")) | |
| ### draw spans result on each page | |
| pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_suff}_spans.pdf")) | |
| ### get markdown content | |
| md_content = pipe_result.get_markdown(image_dir) | |
| ### dump markdown | |
| md_file_path = f"{name_without_suff}.md" | |
| pipe_result.dump_md(md_writer, md_file_path, image_dir) | |
| print(f"Markdown saved to: {os.path.join(local_md_dir, md_file_path)}") | |
| ### get content list content | |
| content_list_content = pipe_result.get_content_list(image_dir) | |
| ### dump content list | |
| pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir) | |
| ### get middle json | |
| middle_json_content = pipe_result.get_middle_json() | |
| ### dump middle json | |
| pipe_result.dump_middle_json(md_writer, f'{name_without_suff}_middle.json') | |
| return md_content | |
| if __name__ == "__main__": | |
| if len(sys.argv) < 2: | |
| print("Usage: python convert_pdf.py <pdf_path>") | |
| sys.exit(1) | |
| pdf_path = sys.argv[1] | |
| if not os.path.exists(pdf_path): | |
| print(f"Error: PDF file not found at {pdf_path}") | |
| sys.exit(1) | |
| convert_pdf(pdf_path) |