mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-10-29 19:07:28 +08:00
76 lines
3.2 KiB
Python
76 lines
3.2 KiB
Python
import pytest
|
|
import os
|
|
from conf import conf
|
|
import os
|
|
import json
|
|
from magic_pdf.pipe.UNIPipe import UNIPipe
|
|
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
|
|
from lib import common
|
|
|
|
pdf_res_path = conf.conf["pdf_res_path"]
|
|
code_path = conf.conf["code_path"]
|
|
pdf_dev_path = conf.conf["pdf_dev_path"]
|
|
class TestCli:
|
|
"""
|
|
test cli
|
|
"""
|
|
def test_pdf_sdk(self):
|
|
"""
|
|
pdf sdk 方式解析
|
|
"""
|
|
demo_names = list()
|
|
pdf_path = os.path.join(pdf_dev_path, "pdf")
|
|
for pdf_file in os.listdir(pdf_path):
|
|
if pdf_file.endswith('.pdf'):
|
|
demo_names.append(pdf_file.split('.')[0])
|
|
for demo_name in demo_names:
|
|
model_path = os.path.join(pdf_dev_path, f"{demo_name}_model.json")
|
|
pdf_path = os.path.join(pdf_dev_path, "pdf", f"{demo_name}.pdf")
|
|
pdf_bytes = open(pdf_path, "rb").read()
|
|
model_json = json.loads(open(model_path, "r", encoding="utf-8").read())
|
|
image_writer = DiskReaderWriter(pdf_dev_path)
|
|
image_dir = str(os.path.basename(pdf_dev_path))
|
|
jso_useful_key = {"_pdf_type": "", "model_list": model_json}
|
|
pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer)
|
|
pipe.pipe_classify()
|
|
pipe.pipe_parse()
|
|
md_content = pipe.pipe_mk_markdown(image_dir, drop_mode="none")
|
|
dir_path = os.path.join(pdf_dev_path, "mineru")
|
|
if not os.path.exists(dir_path):
|
|
os.makedirs(dir_path, exist_ok=True)
|
|
res_path = os.path.join(dir_path, f"{demo_name}.md")
|
|
with open(res_path, "w+", encoding="utf-8") as f:
|
|
f.write(md_content)
|
|
common.count_folders_and_check_contents(res_path)
|
|
|
|
# def test_pdf_specify_jsonl(self):
|
|
# """
|
|
# 输入jsonl, 默认方式解析
|
|
# """
|
|
# cmd = "cd %s && export PYTHONPATH=. && python magic_pdf/cli/magicpdf.py json-command --json 's3://llm-process-pperf/ebook_index_textbook_40k/中高考&竞赛知识点/part-663f1ef5e7c1-009416.jsonl?bytes=0,1133972'" % (code_path)
|
|
# logging.info(cmd)
|
|
# common.check_shell(cmd)
|
|
# #common.count_folders_and_check_contents(pdf_res_path)
|
|
|
|
# def test_pdf_specify_jsonl_txt(self):
|
|
# """
|
|
# 输入jsonl, txt方式解析
|
|
# """
|
|
# cmd = "cd %s && export PYTHONPATH=. && python magic_pdf/cli/magicpdf.py json-command --json 's3://llm-process-pperf/ebook_index_textbook_40k/中高考&竞赛知识点/part-663f1ef5e7c1-009416.jsonl?bytes=0,1133972' --method txt" % (code_path)
|
|
# logging.info(cmd)
|
|
# common.check_shell(cmd)
|
|
# #common.count_folders_and_check_contents(pdf_res_path)
|
|
#
|
|
# def test_pdf_specify_jsonl_ocr(self):
|
|
# """
|
|
# 输入jsonl, ocr方式解析
|
|
# """
|
|
# cmd = "cd %s && export PYTHONPATH=. && python magic_pdf/cli/magicpdf.py json-command --json 's3://llm-process-pperf/ebook_index_textbook_40k/中高考&竞赛知识点/part-663f1ef5e7c1-009416.jsonl?bytes=0,1133972' --method ocr" % (code_path)
|
|
# logging.info(cmd)
|
|
# common.check_shell(cmd)
|
|
# #common.count_folders_and_check_contents(pdf_res_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|