MinerU/demo/ocr_demo.py

import json
import os

from loguru import logger
from pathlib import Path

from magic_pdf.pipeline_ocr import ocr_parse_pdf_core
from magic_pdf.spark.s3 import get_s3_config
from demo.demo_commons import get_json_from_local_or_s3
from magic_pdf.dict2md.ocr_mkcontent import (
    ocr_mk_mm_markdown_with_para,
    make_standard_format_with_para
)
from magic_pdf.libs.commons import join_path, read_file


def save_markdown(markdown_text, input_filepath):
    # 获取输入文件的目录
    directory = os.path.dirname(input_filepath)
    # 获取输入文件的文件名（不带扩展名）
    base_name = os.path.basename(input_filepath)
    file_name_without_ext = os.path.splitext(base_name)[0]
    # 定义输出文件的路径
    output_filepath = os.path.join(directory, f"{file_name_without_ext}.md")

    # 将Markdown文本写入.md文件
    with open(output_filepath, 'w', encoding='utf-8') as file:
        file.write(markdown_text)


def read_json_file(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    return data


def ocr_local_parse(ocr_pdf_path, ocr_json_file_path):
    try:
        ocr_pdf_model_info = read_json_file(ocr_json_file_path)
        pth = Path(ocr_json_file_path)
        book_name = pth.name
        pdf_bytes = read_file(ocr_pdf_path, None)
        ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info)
    except Exception as e:
        logger.exception(e)


def ocr_online_parse(book_name, start_page_id=0, debug_mode=True):
    try:
        json_object = get_json_from_local_or_s3(book_name)
        # logger.info(json_object)
        s3_pdf_path = json_object["file_location"]
        s3_config = get_s3_config(s3_pdf_path)
        pdf_bytes = read_file(s3_pdf_path, s3_config)
        ocr_pdf_model_info = json_object.get("doc_layout_result")
        ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info)
    except Exception as e:
        logger.exception(e)


def ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info, start_page_id=0):
    save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
    save_path = join_path(save_tmp_path, "md")
    save_path_with_bookname = os.path.join(save_path, book_name)
    text_content_save_path = f"{save_path_with_bookname}/book.md"
    pdf_info_dict, parse_time = ocr_parse_pdf_core(pdf_bytes, ocr_pdf_model_info, book_name, start_page_id=start_page_id, debug_mode=True)

    parent_dir = os.path.dirname(text_content_save_path)
    if not os.path.exists(parent_dir):
        os.makedirs(parent_dir)

    # markdown_content = mk_nlp_markdown(pdf_info_dict)
    markdown_content = ocr_mk_mm_markdown_with_para(pdf_info_dict)
    # markdown_pagination = ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict)

    with open(text_content_save_path, "w", encoding="utf-8") as f:
        f.write(markdown_content)

    standard_format = make_standard_format_with_para(pdf_info_dict)
    standard_format_save_path = f"{save_path_with_bookname}/standard_format.txt"
    with open(standard_format_save_path, "w", encoding="utf-8") as f:
        # 将standard_format dump成json文本并保存
        f.write(json.dumps(standard_format, ensure_ascii=False))


if __name__ == '__main__':
    pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.pdf"
    json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.json"
    # ocr_local_parse(pdf_path, json_file_path)
    book_name = "科数网/edu_00011318"
    ocr_online_parse(book_name)
    
    pass