opensource/MinerU
3
0
Fork 0
mirror of https://github.com/opendatalab/MinerU.git synced 2026-09-01 04:26:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

update test case

Browse Source
This commit is contained in:
quyuan
2024-11-17 17:19:48 +08:00
parent 12ced4f39e
commit 59ad273784
3 changed files with 113 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_cli/conftest.py
-7
View File
@@ -8,10 +8,3 @@ def clear_gpu_memory():
torch.cuda.empty_cache()
print("GPU memory cleared.")
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_teardown(item, nextitem):
'''
clear GPU memory after each test
'''
yield
clear_gpu_memory()
tests/test_cli/lib/common.py
+15 -4
View File
@@ -3,6 +3,15 @@ import os
import shutil
import re
import json
import torch
def clear_gpu_memory():
'''
clear GPU memory
'''
torch.cuda.empty_cache()
print("GPU memory cleared.")
def check_shell(cmd):
"""shell successful."""
res = os.system(cmd)
@@ -10,11 +19,12 @@ def check_shell(cmd):
def update_config_file(file_path, key, value):
"""update config file."""
with open(file_path, 'r', encoding="utf-8") as f:
config = json.loads(f.read())
with open(file_path, 'r', encoding="utf-8") as fr:
config = json.loads(fr.read())
config[key] = value
with open(file_path, 'w', encoding="utf-8") as f:
f.write(json.dumps(config))
# 保存修改后的内容
with open(file_path, 'w', encoding='utf-8') as fw:
json.dump(config, fw, ensure_ascii=False, indent=4)
def cli_count_folders_and_check_contents(file_path):
"""" count cli files."""
@@ -33,6 +43,7 @@ def sdk_count_folders_and_check_contents(file_path):
exit(1)
def delete_file(path):
"""delete file."""
if not os.path.exists(path):
tests/test_cli/test_cli_sdk.py
+98 -14
View File
@@ -13,10 +13,19 @@ model_config.__use_inside_model__ = True
pdf_res_path = conf.conf['pdf_res_path']
code_path = conf.conf['code_path']
pdf_dev_path = conf.conf['pdf_dev_path']
magic_pdf_config = "/home/quyuan/magic-pdf.json"
class TestCli:
"""test cli."""
@pytest.fixture(autouse=True)
def setup(self):
"""
init
"""
common.clear_gpu_memory()
common.update_config_file(magic_pdf_config, "device-mode", "cuda")
# 这里可以添加任何前置操作
yield
@pytest.mark.P0
def test_pdf_auto_sdk(self):
@@ -291,22 +300,32 @@ class TestCli:
def test_local_magic_pdf_open_st_table(self):
"""magic pdf cli open st table."""
time.sleep(2)
pre_cmd = "cp ~/magic_pdf_st.json ~/magic-pdf.json"
print (pre_cmd)
os.system(pre_cmd)
#pre_cmd = "cp ~/magic_pdf_st.json ~/magic-pdf.json"
value = {
"model": "struct_eqtable",
"enable": True,
"max_time": 400
}
common.update_config_file(magic_pdf_config, "table-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
os.system(cli_cmd)
res = common.check_latex_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
assert res is True
@pytest.mark.P1
def test_local_magic_pdf_open_html_table(self):
"""magic pdf cli open html table."""
def test_local_magic_pdf_open_tablemaster_cuda(self):
"""magic pdf cli open table master html table cuda mode."""
time.sleep(2)
pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
os.system(pre_cmd)
#pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
#os.system(pre_cmd)
value = {
"model": "tablemaster",
"enable": True,
"max_time": 400
}
common.update_config_file(magic_pdf_config, "table-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
@@ -315,11 +334,69 @@ class TestCli:
assert res is True
@pytest.mark.P1
def test_magic_pdf_close_html_table_cpu(self):
def test_local_magic_pdf_open_rapidai_table(self):
"""magic pdf cli open rapid ai table."""
time.sleep(2)
#pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
#os.system(pre_cmd)
value = {
"model": "rapid_table",
"enable": True,
"max_time": 400
}
common.update_config_file(magic_pdf_config, "table-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
os.system(cli_cmd)
res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
assert res is True
@pytest.mark.P1
def test_local_magic_pdf_doclayout_yolo(self):
"""magic pdf cli open doclyaout yolo."""
time.sleep(2)
#pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
#os.system(pre_cmd)
value = {
"model": "doclayout_yolo"
}
common.update_config_file(magic_pdf_config, "layout-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
os.system(cli_cmd)
common.cli_count_folders_and_check_contents(os.path.join(pdf_res_path, "test_rearch_report", "auto"))
@pytest.mark.P1
def test_local_magic_pdf_layoutlmv3_yolo(self):
"""magic pdf cli open layoutlmv3."""
time.sleep(2)
value = {
"model": "layoutlmv3"
}
common.update_config_file(magic_pdf_config, "layout-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
os.system(cli_cmd)
common.cli_count_folders_and_check_contents(os.path.join(pdf_res_path, "test_rearch_report", "auto"))
#res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
@pytest.mark.P1
def test_magic_pdf_open_tablemaster_table_cpu(self):
"""magic pdf cli close html table cpu mode."""
time.sleep(2)
pre_cmd = "cp ~/magic_pdf_html_table_cpu.json ~/magic-pdf.json"
os.system(pre_cmd)
#pre_cmd = "cp ~/magic_pdf_html_table_cpu.json ~/magic-pdf.json"
#os.system(pre_cmd)
value = {
"model": "tablemaster",
"enable": True,
"max_time": 400
}
common.update_config_file(magic_pdf_config, "table-config", value)
common.update_config_file(magic_pdf_config, "device-mode", "cpu")
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
@@ -327,12 +404,19 @@ class TestCli:
res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
assert res is True
@pytest.mark.P1
def test_local_magic_pdf_close_html_table(self):
"""magic pdf cli close table."""
time.sleep(2)
pre_cmd = "cp ~/magic_pdf_close_table.json ~/magic-pdf.json"
os.system(pre_cmd)
#pre_cmd = "cp ~/magic_pdf_close_table.json ~/magic-pdf.json"
#os.system(pre_cmd)
value = {
"model": "tablemaster",
"enable": False,
"max_time": 400
}
common.update_config_file(magic_pdf_config, "table-config", value)
pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
common.delete_file(pdf_res_path)
cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)