-
Notifications
You must be signed in to change notification settings - Fork 2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(demo): add demo script for PDF processing
- Create demo.py script for PDF file processing - Implement PDF reading, classification, and inference usingOpendatalab's magic_pdf library- Add pipelines for OCR and text modes - Include result visualization and markdown export
- Loading branch information
Showing
1 changed file
with
58 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,58 @@ | ||
# Copyright (c) Opendatalab. All rights reserved. | ||
import os | ||
|
||
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 | ||
|
||
# args | ||
pdf_file_name = "demo1.pdf" # replace with the real pdf path | ||
name_without_suff = pdf_file_name.split(".")[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) | ||
|
||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( | ||
local_md_dir | ||
) | ||
image_dir = str(os.path.basename(local_image_dir)) | ||
|
||
# read bytes | ||
reader1 = FileBasedDataReader("") | ||
pdf_bytes = reader1.read(pdf_file_name) # read the pdf content | ||
|
||
# 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")) | ||
|
||
### 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")) | ||
|
||
### dump markdown | ||
pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir) | ||
|
||
### dump content list | ||
pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir) |