wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Waiting to run
Main / Python 3.11 - Build (push) Waiting to run
Main / Python 3.11 - Lint (push) Waiting to run
Main / Python 3.11 - Style (push) Waiting to run
Main / Python 3.11 - Test (push) Waiting to run
Main / GPU CI (push) Blocked by required conditions
Main / Release (push) Blocked by required conditions
Main / Build and Push Docker Images (push) Blocked by required conditions
75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
import os
|
|
import tempfile
|
|
|
|
from magic_pdf.config.enums import SupportedPdfParseMethod
|
|
from magic_pdf.data.data_reader_writer import FileBasedDataReader, FileBasedDataWriter
|
|
from magic_pdf.data.dataset import PymuDocDataset
|
|
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
|
|
from pypdf import PdfReader, PdfWriter
|
|
|
|
|
|
def run_mineru(pdf_path: str, page_num: int = 1) -> str:
|
|
output_folder = tempfile.TemporaryDirectory()
|
|
image_output_folder = tempfile.TemporaryDirectory()
|
|
|
|
# Initialize writers (same for all PDFs)
|
|
image_writer = FileBasedDataWriter(image_output_folder.name)
|
|
md_writer = FileBasedDataWriter(output_folder.name)
|
|
|
|
if page_num > 0: # If a specific page is requested
|
|
reader = PdfReader(pdf_path)
|
|
|
|
# Check if the requested page exists
|
|
if page_num > len(reader.pages):
|
|
raise ValueError(f"Page {page_num} does not exist in the PDF. PDF has {len(reader.pages)} pages.")
|
|
|
|
# Create a new PDF with just the requested page
|
|
writer = PdfWriter()
|
|
# pypdf uses 0-based indexing, so subtract 1 from page_num
|
|
writer.add_page(reader.pages[page_num - 1])
|
|
|
|
# Save the extracted page to a temporary file
|
|
temp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
|
|
temp_file.close() # Close the file but keep the name
|
|
|
|
with open(temp_file.name, "wb") as output_pdf:
|
|
writer.write(output_pdf)
|
|
|
|
pdf_to_process = temp_file.name
|
|
else:
|
|
pdf_to_process = pdf_path
|
|
|
|
try:
|
|
# Read the PDF file bytes
|
|
reader = FileBasedDataReader("")
|
|
pdf_bytes = reader.read(pdf_to_process)
|
|
|
|
# Create dataset instance
|
|
ds = PymuDocDataset(pdf_bytes)
|
|
|
|
# Inference: decide whether to run OCR mode based on dataset classification
|
|
if ds.classify() == SupportedPdfParseMethod.OCR:
|
|
infer_result = ds.apply(doc_analyze, ocr=True)
|
|
pipe_result = infer_result.pipe_ocr_mode(image_writer)
|
|
else:
|
|
infer_result = ds.apply(doc_analyze, ocr=False)
|
|
pipe_result = infer_result.pipe_txt_mode(image_writer)
|
|
|
|
# Generate markdown content; the image directory is the basename of the images output folder
|
|
image_dir_basename = os.path.basename(image_output_folder.name)
|
|
# md_content = pipe_result.get_markdown(image_dir_basename)
|
|
|
|
# Dump markdown file
|
|
with tempfile.NamedTemporaryFile("w+", suffix="md") as tf:
|
|
pipe_result.dump_md(md_writer, tf.name, image_dir_basename)
|
|
tf.flush()
|
|
|
|
tf.seek(0)
|
|
md_data = tf.read()
|
|
|
|
return md_data
|
|
finally:
|
|
# Clean up the temporary file if it was created
|
|
if temp_file and os.path.exists(temp_file.name):
|
|
os.unlink(temp_file.name)
|