mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-09-01 15:57:25 +08:00
Xiaomeng Zhao
+14
-10
@@ -14,27 +14,31 @@ on:
|
||||
jobs:
|
||||
cli-test:
|
||||
if: github.repository == 'opendatalab/MinerU'
|
||||
runs-on: pdf
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 240
|
||||
strategy:
|
||||
fail-fast: true
|
||||
|
||||
steps:
|
||||
- name: PDF cli
|
||||
uses: actions/checkout@v3
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: dev
|
||||
fetch-depth: 2
|
||||
|
||||
- name: install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
|
||||
- name: install&test
|
||||
run: |
|
||||
source activate mineru
|
||||
conda env list
|
||||
pip show coverage
|
||||
cd $GITHUB_WORKSPACE && sh tests/retry_env.sh
|
||||
# cd $GITHUB_WORKSPACE && python tests/clean_coverage.py
|
||||
# cd $GITHUB_WORKSPACE && coverage run -m pytest tests/unittest/ --cov=magic_pdf/ --cov-report html --cov-report term-missing
|
||||
# cd $GITHUB_WORKSPACE && python tests/get_coverage.py
|
||||
cd $GITHUB_WORKSPACE && pytest -m P0 -s -v tests/test_cli/test_cli_sdk.py
|
||||
uv --version
|
||||
uv venv --python 3.12
|
||||
source .venv/bin/activate
|
||||
uv pip install .[test]
|
||||
cd $GITHUB_WORKSPACE && python tests/clean_coverage.py
|
||||
cd $GITHUB_WORKSPACE && coverage run
|
||||
cd $GITHUB_WORKSPACE && python tests/get_coverage.py
|
||||
|
||||
|
||||
notify_to_feishu:
|
||||
if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure')}}
|
||||
|
||||
@@ -12,33 +12,36 @@ on:
|
||||
- "**.md"
|
||||
jobs:
|
||||
cli-test:
|
||||
if: github.repository == 'opendatalab/MinerU'
|
||||
runs-on: pdf
|
||||
# if: github.repository == 'opendatalab/MinerU'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 240
|
||||
strategy:
|
||||
fail-fast: true
|
||||
|
||||
steps:
|
||||
- name: PDF cli
|
||||
uses: actions/checkout@v3
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: dev
|
||||
fetch-depth: 2
|
||||
|
||||
- name: install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
|
||||
- name: install&test
|
||||
run: |
|
||||
source activate mineru
|
||||
conda env list
|
||||
pip show coverage
|
||||
cd $GITHUB_WORKSPACE && sh tests/retry_env.sh
|
||||
# cd $GITHUB_WORKSPACE && python tests/clean_coverage.py
|
||||
# cd $GITHUB_WORKSPACE && coverage run -m pytest tests/unittest/ --cov=magic_pdf/ --cov-report html --cov-report term-missing
|
||||
# cd $GITHUB_WORKSPACE && python tests/get_coverage.py
|
||||
cd $GITHUB_WORKSPACE && pytest -s -v tests/test_cli/test_cli_sdk.py
|
||||
uv --version
|
||||
uv venv --python 3.12
|
||||
source .venv/bin/activate
|
||||
uv pip install .[test]
|
||||
cd $GITHUB_WORKSPACE && python tests/clean_coverage.py
|
||||
cd $GITHUB_WORKSPACE && coverage run
|
||||
cd $GITHUB_WORKSPACE && python tests/get_coverage.py
|
||||
|
||||
notify_to_feishu:
|
||||
if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure')}}
|
||||
# if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure')}}
|
||||
needs: cli-test
|
||||
runs-on: pdf
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: get_actor
|
||||
run: |
|
||||
@@ -57,5 +60,5 @@ jobs:
|
||||
|
||||
- name: notify
|
||||
run: |
|
||||
#echo ${{ secrets.USER_ID }}
|
||||
curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'$USER_ID'"}]]}}}}' $WEBHOOK_URL
|
||||
echo ${{ secrets.USER_ID }}
|
||||
curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"}]]}}}}' $WEBHOOK_URL
|
||||
|
||||
@@ -100,6 +100,9 @@ plugins:
|
||||
- search
|
||||
- i18n:
|
||||
docs_structure: folder
|
||||
fallback_to_default: true
|
||||
reconfigure_material: true
|
||||
reconfigure_search: true
|
||||
languages:
|
||||
- locale: en
|
||||
default: true
|
||||
|
||||
+45
-2
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
|
||||
[project]
|
||||
name = "mineru"
|
||||
dynamic = ["version"]
|
||||
license = {text = "AGPL-3.0"}
|
||||
license = { text = "AGPL-3.0" }
|
||||
description = "A practical tool for converting PDF to Markdown"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10,<3.14"
|
||||
@@ -38,6 +38,14 @@ dependencies = [
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
test = [
|
||||
"mineru[core]",
|
||||
"pytest",
|
||||
"pytest-cov",
|
||||
"coverage",
|
||||
"beautifulsoup4",
|
||||
"fuzzywuzzy"
|
||||
]
|
||||
vlm = [
|
||||
"transformers>=4.51.1",
|
||||
"torch>=2.6.0",
|
||||
@@ -112,7 +120,7 @@ mineru-api = "mineru.cli.fast_api:main"
|
||||
mineru-gradio = "mineru.cli.gradio_app:main"
|
||||
|
||||
[tool.setuptools.dynamic]
|
||||
version = {attr = "mineru.version.__version__"}
|
||||
version = { attr = "mineru.version.__version__" }
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
include = ["mineru*"]
|
||||
@@ -125,3 +133,38 @@ namespaces = false
|
||||
[tool.setuptools]
|
||||
include-package-data = true
|
||||
zip-safe = false
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
addopts = "-s --cov=mineru --cov-report html"
|
||||
|
||||
[tool.coverage.run]
|
||||
command_line = "-m pytest tests/unittest/test_e2e.py"
|
||||
source = ["mineru/"]
|
||||
omit = [
|
||||
"*/vlm_sglang_model/*",
|
||||
"*/gradio_app.py",
|
||||
"*/models_download.py",
|
||||
"*/fast_api.py",
|
||||
"*/cli/client.py",
|
||||
"*/sglang_engine_predictor.py",
|
||||
"*/vlm_sglang_server.py",
|
||||
"*/cli_parser.py",
|
||||
"*/run_async.py"
|
||||
]
|
||||
|
||||
[tool.coverage.html]
|
||||
directory = "htmlcov"
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_also = [
|
||||
'def __repr__',
|
||||
'if self.debug:',
|
||||
'if settings.DEBUG',
|
||||
'raise AssertionError',
|
||||
'raise NotImplementedError',
|
||||
'if 0:',
|
||||
'if __name__ == .__main__.:',
|
||||
'if TYPE_CHECKING:',
|
||||
'class .*\bProtocol\):',
|
||||
'@(abc\.)?abstractmethod',
|
||||
]
|
||||
@@ -1,28 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
max_retries=5
|
||||
retry_count=0
|
||||
|
||||
while true; do
|
||||
# prepare env
|
||||
#python -m pip install -r requirements-qa.txt
|
||||
#python -m pip install -U magic-pdf[full] --extra-index-url https://wheels.myhloli.com -i https://mirrors.aliyun.com/pypi/simple
|
||||
pip install -e .
|
||||
python -m pip install paddlepaddle-gpu==3.0.0b1 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
|
||||
pip install modelscope
|
||||
wget https://gcore.jsdelivr.net/gh/opendatalab/MinerU@master/scripts/download_models.py -O download_models.py
|
||||
python download_models.py
|
||||
exit_code=$?
|
||||
if [ $exit_code -eq 0 ]; then
|
||||
echo "test.sh 成功执行！"
|
||||
break
|
||||
else
|
||||
let retry_count+=1
|
||||
if [ $retry_count -ge $max_retries ]; then
|
||||
echo "达到最大重试次数 ($max_retries)，放弃重试。"
|
||||
exit 1
|
||||
fi
|
||||
echo "test.sh 执行失败 (退出码: $exit_code)。尝试第 $retry_count 次重试..."
|
||||
sleep 5
|
||||
fi
|
||||
done
|
||||
@@ -1,10 +0,0 @@
|
||||
import os
|
||||
conf = {
|
||||
"code_path": os.environ.get('GITHUB_WORKSPACE'),
|
||||
"pdf_dev_path" : os.environ.get('GITHUB_WORKSPACE') + "/tests/test_cli/pdf_dev",
|
||||
#"code_path": "/home/quyuan/ci/actions-runner/MinerU",
|
||||
#"pdf_dev_path": "/home/quyuan/ci/actions-runner/MinerU/tests/test_cli/pdf_dev",
|
||||
"pdf_res_path": "/tmp/magic-pdf",
|
||||
"jsonl_path": "s3://llm-qatest-pnorm/mineru/test/line1.jsonl",
|
||||
"s3_pdf_path": "s3://llm-qatest-pnorm/mineru/test/test_rearch_report.pdf"
|
||||
}
|
||||
@@ -1,10 +0,0 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
def clear_gpu_memory():
|
||||
'''
|
||||
clear GPU memory
|
||||
'''
|
||||
torch.cuda.empty_cache()
|
||||
print("GPU memory cleared.")
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
"""
|
||||
calculate_score
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
from Levenshtein import distance
|
||||
from lib import scoring
|
||||
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
|
||||
from nltk.tokenize import word_tokenize
|
||||
import nltk
|
||||
nltk.download('punkt')
|
||||
|
||||
class Scoring:
|
||||
"""
|
||||
calculate_score
|
||||
"""
|
||||
def __init__(self, result_path):
|
||||
"""
|
||||
init
|
||||
"""
|
||||
self.edit_distances = []
|
||||
self.bleu_scores = []
|
||||
self.sim_scores = []
|
||||
self.filenames = []
|
||||
self.score_dict = {}
|
||||
self.anntion_cnt = 0
|
||||
self.fw = open(result_path, "w+", encoding='utf-8')
|
||||
|
||||
def simple_bleu_score(self, candidate, reference):
|
||||
"""
|
||||
get bleu score
|
||||
"""
|
||||
candidate_tokens = word_tokenize(candidate)
|
||||
reference_tokens = word_tokenize(reference)
|
||||
return sentence_bleu([reference_tokens], candidate_tokens, smoothing_function=SmoothingFunction().method1)
|
||||
|
||||
|
||||
def preprocess_string(self, s):
|
||||
"""
|
||||
preprocess_string
|
||||
"""
|
||||
sub_enter = re.sub(r'\n+', '\n', s)
|
||||
return re.sub(r' ', ' ', sub_enter)
|
||||
|
||||
def calculate_similarity(self, annotion, actual, tool_type):
|
||||
"""
|
||||
calculate_similarity
|
||||
"""
|
||||
class_dict = {}
|
||||
edit_distances = []
|
||||
bleu_scores = []
|
||||
sim_scores = list()
|
||||
total_file = 0
|
||||
for filename in os.listdir(annotion):
|
||||
if filename.endswith('.md') and not filename.startswith('.'):
|
||||
total_file = total_file + 1
|
||||
with open(os.path.join(annotion, filename), 'r', encoding='utf-8') as file_a:
|
||||
content_a = file_a.read()
|
||||
self.anntion_cnt = self.anntion_cnt + 1
|
||||
filepath_b = os.path.join(actual, filename)
|
||||
if os.path.exists(filepath_b):
|
||||
with open(filepath_b, 'r', encoding='utf-8') as file_b:
|
||||
content_b = file_b.read()
|
||||
self.filenames.append(filename)
|
||||
edit_dist = distance(self.preprocess_string(content_b),self.preprocess_string(content_a)) / max(len(content_a), len(content_b))
|
||||
self.edit_distances.append(edit_dist)
|
||||
edit_distances.append(edit_dist)
|
||||
bleu_score = self.simple_bleu_score(content_b, content_a)
|
||||
bleu_scores.append(bleu_score)
|
||||
self.bleu_scores.append(bleu_score)
|
||||
score = scoring.score_text(content_b, content_a)
|
||||
sim_scores.append(score)
|
||||
self.sim_scores.append(score)
|
||||
class_dict[filename] = {"edit_dist": edit_dist, "bleu_score": bleu_score, "sim_score": score}
|
||||
self.score_dict[filename] = {"edit_dist": edit_dist, "bleu_score": bleu_score, "sim_score": score}
|
||||
else:
|
||||
print(f"File {filename} not found in actual directory.")
|
||||
class_average_edit_distance = sum(edit_distances) / len(edit_distances) if edit_distances else 0
|
||||
class_average_bleu_score = sum(bleu_scores) / len(bleu_scores) if bleu_scores else 0
|
||||
class_average_sim_score = sum(sim_scores) / len(sim_scores) if sim_scores else 0
|
||||
self.fw.write(json.dumps(class_dict, ensure_ascii=False) + "\n")
|
||||
ratio = len(class_dict)/total_file
|
||||
self.fw.write(f"{tool_type} extract ratio: {ratio}" + "\n")
|
||||
self.fw.write(f"{tool_type} Average Levenshtein Distance: {class_average_edit_distance}" + "\n")
|
||||
self.fw.write(f"{tool_type} Average BLEU Score: {class_average_bleu_score}" + "\n")
|
||||
self.fw.write(f"{tool_type} Average Sim Score: {class_average_sim_score}" + "\n")
|
||||
print (f"{tool_type} extract ratio: {ratio}")
|
||||
print (f"{tool_type} Average Levenshtein Distance: {class_average_edit_distance}")
|
||||
print (f"{tool_type} Average BLEU Score: {class_average_bleu_score}")
|
||||
print (f"{tool_type} Average Sim Score: {class_average_sim_score}")
|
||||
return self.score_dict
|
||||
|
||||
def summary_scores(self):
|
||||
"""
|
||||
calculate the average of edit distance, bleu score and sim score
|
||||
"""
|
||||
over_all_dict = dict()
|
||||
average_edit_distance = sum(self.edit_distances) / len(self.edit_distances) if self.edit_distances else 0
|
||||
average_bleu_score = sum(self.bleu_scores) / len(self.bleu_scores) if self.bleu_scores else 0
|
||||
average_sim_score = sum(self.sim_scores) / len(self.sim_scores) if self.sim_scores else 0
|
||||
over_all_dict["average_edit_distance"] = average_edit_distance
|
||||
over_all_dict["average_bleu_score"] = average_bleu_score
|
||||
over_all_dict["average_sim_score"] = average_sim_score
|
||||
self.fw.write(json.dumps(over_all_dict, ensure_ascii=False) + "\n")
|
||||
return over_all_dict
|
||||
|
||||
def calculate_similarity_total(self, tool_type, download_dir):
|
||||
"""
|
||||
calculate the average of edit distance, bleu score and sim score
|
||||
"""
|
||||
annotion = os.path.join(download_dir, "annotations", "cleaned")
|
||||
actual = os.path.join(download_dir, tool_type, "cleaned")
|
||||
score = self.calculate_similarity(annotion, actual, tool_type)
|
||||
return score
|
||||
|
||||
@@ -1,90 +0,0 @@
|
||||
"""common definitions."""
|
||||
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)
|
||||
assert res == 0
|
||||
|
||||
def update_config_file(file_path, key, value):
|
||||
"""update config file."""
|
||||
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 fw:
|
||||
json.dump(config, fw, ensure_ascii=False, indent=4)
|
||||
|
||||
def cli_count_folders_and_check_contents(file_path):
|
||||
"""" count cli files."""
|
||||
if os.path.exists(file_path):
|
||||
for files in os.listdir(file_path):
|
||||
folder_count = os.path.getsize(os.path.join(file_path, files))
|
||||
assert folder_count > 0
|
||||
assert len(os.listdir(file_path)) > 5
|
||||
|
||||
def sdk_count_folders_and_check_contents(file_path):
|
||||
"""count folders."""
|
||||
if os.path.exists(file_path):
|
||||
file_count = os.path.getsize(file_path)
|
||||
assert file_count > 0
|
||||
else:
|
||||
exit(1)
|
||||
|
||||
|
||||
|
||||
def delete_file(path):
|
||||
"""delete file."""
|
||||
if not os.path.exists(path):
|
||||
if os.path.isfile(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
print(f"File '{path}' deleted.")
|
||||
except TypeError as e:
|
||||
print(f"Error deleting file '{path}': {e}")
|
||||
elif os.path.isdir(path):
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
print(f"Directory '{path}' and its contents deleted.")
|
||||
except TypeError as e:
|
||||
print(f"Error deleting directory '{path}': {e}")
|
||||
|
||||
def check_latex_table_exists(file_path):
|
||||
"""check latex table exists."""
|
||||
pattern = r'\\begin\{tabular\}.*?\\end\{tabular\}'
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
matches = re.findall(pattern, content, re.DOTALL)
|
||||
return len(matches) > 0
|
||||
|
||||
def check_html_table_exists(file_path):
|
||||
"""check html table exists."""
|
||||
pattern = r'<table.*?>.*?</table>'
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
matches = re.findall(pattern, content, re.DOTALL)
|
||||
return len(matches) > 0
|
||||
|
||||
def check_close_tables(file_path):
|
||||
"""delete no tables."""
|
||||
latex_pattern = r'\\begin\{tabular\}.*?\\end\{tabular\}'
|
||||
html_pattern = r'<table.*?>.*?</table>'
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
latex_matches = re.findall(latex_pattern, content, re.DOTALL)
|
||||
html_matches = re.findall(html_pattern, content, re.DOTALL)
|
||||
if len(latex_matches) == 0 and len(html_matches) == 0:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
@@ -1,128 +0,0 @@
|
||||
"""
|
||||
clean data
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import htmltabletomd # type: ignore
|
||||
import pypandoc
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="get tool type")
|
||||
parser.add_argument(
|
||||
"--tool_name",
|
||||
type=str,
|
||||
required=True,
|
||||
help="input tool name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--download_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="input download dir",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
def clean_markdown_images(content):
|
||||
"""
|
||||
clean markdown images
|
||||
"""
|
||||
pattern = re.compile(r'!\[[^\]]*\]\([^)]*\)', re.IGNORECASE)
|
||||
cleaned_content = pattern.sub('', content)
|
||||
return cleaned_content
|
||||
|
||||
def clean_ocrmath_photo(content):
|
||||
"""
|
||||
clean ocrmath photo
|
||||
"""
|
||||
pattern = re.compile(r'\\includegraphics\[.*?\]\{.*?\}', re.IGNORECASE)
|
||||
cleaned_content = pattern.sub('', content)
|
||||
return cleaned_content
|
||||
|
||||
def convert_html_table_to_md(html_table):
|
||||
"""
|
||||
convert html table to markdown table
|
||||
"""
|
||||
lines = html_table.strip().split('\n')
|
||||
md_table = ''
|
||||
if lines and '<tr>' in lines[0]:
|
||||
in_thead = True
|
||||
for line in lines:
|
||||
if '<th>' in line:
|
||||
cells = re.findall(r'<th>(.*?)</th>', line)
|
||||
md_table += '| ' + ' | '.join(cells) + ' |\n'
|
||||
in_thead = False
|
||||
elif '<td>' in line and not in_thead:
|
||||
cells = re.findall(r'<td>(.*?)</td>', line)
|
||||
md_table += '| ' + ' | '.join(cells) + ' |\n'
|
||||
md_table = md_table.rstrip() + '\n'
|
||||
return md_table
|
||||
|
||||
def convert_latext_to_md(content):
|
||||
"""
|
||||
convert latex table to markdown table
|
||||
"""
|
||||
tables = re.findall(r'\\begin\{tabular\}(.*?)\\end\{tabular\}', content, re.DOTALL)
|
||||
placeholders = []
|
||||
for table in tables:
|
||||
placeholder = f"<!-- TABLE_PLACEHOLDER_{len(placeholders)} -->"
|
||||
replace_str = f"\\begin{{tabular}}{table}cl\\end{{tabular}}"
|
||||
content = content.replace(replace_str, placeholder)
|
||||
try:
|
||||
pypandoc.convert_text(replace_str, format="latex", to="md", outputfile="output.md", encoding="utf-8")
|
||||
except:
|
||||
markdown_string = replace_str
|
||||
else:
|
||||
markdown_string = open('output.md', 'r', encoding='utf-8').read()
|
||||
placeholders.append((placeholder, markdown_string))
|
||||
new_content = content
|
||||
for placeholder, md_table in placeholders:
|
||||
new_content = new_content.replace(placeholder, md_table)
|
||||
# 写入文件
|
||||
return new_content
|
||||
|
||||
|
||||
def convert_htmltale_to_md(content):
|
||||
"""
|
||||
convert html table to markdown table
|
||||
"""
|
||||
tables = re.findall(r'<table>(.*?)</table>', content, re.DOTALL)
|
||||
placeholders = []
|
||||
for table in tables:
|
||||
placeholder = f"<!-- TABLE_PLACEHOLDER_{len(placeholders)} -->"
|
||||
content = content.replace(f"<table>{table}</table>", placeholder)
|
||||
try:
|
||||
convert_table = htmltabletomd.convert_table(table)
|
||||
except:
|
||||
convert_table = table
|
||||
placeholders.append((placeholder,convert_table))
|
||||
new_content = content
|
||||
for placeholder, md_table in placeholders:
|
||||
new_content = new_content.replace(placeholder, md_table)
|
||||
# 写入文件
|
||||
return new_content
|
||||
|
||||
def clean_data(prod_type, download_dir):
|
||||
"""
|
||||
clean data
|
||||
"""
|
||||
tgt_dir = os.path.join(download_dir, prod_type, "cleaned")
|
||||
if not os.path.exists(tgt_dir):
|
||||
os.makedirs(tgt_dir)
|
||||
source_dir = os.path.join(download_dir, prod_type)
|
||||
filenames = os.listdir(source_dir)
|
||||
for filename in filenames:
|
||||
if filename.endswith('.md'):
|
||||
input_file = os.path.join(source_dir, filename)
|
||||
output_file = os.path.join(tgt_dir, "cleaned_" + filename)
|
||||
with open(input_file, 'r', encoding='utf-8') as fr:
|
||||
content = fr.read()
|
||||
new_content = clean_markdown_images(content)
|
||||
with open(output_file, 'w', encoding='utf-8') as fw:
|
||||
fw.write(new_content)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
tool_type = args.tool_name
|
||||
download_dir = args.download_dir
|
||||
clean_data(tool_type, download_dir)
|
||||
@@ -1,51 +0,0 @@
|
||||
"""
|
||||
Calculate simscore, refer to (https://github.com/VikParuchuri/marker?tab=readme-ov-file)
|
||||
"""
|
||||
import math
|
||||
|
||||
from rapidfuzz import fuzz
|
||||
import re
|
||||
import regex
|
||||
from statistics import mean
|
||||
|
||||
CHUNK_MIN_CHARS = 25
|
||||
|
||||
def chunk_text(text, chunk_len=500):
|
||||
chunks = [text[i:i+chunk_len] for i in range(0, len(text), chunk_len)]
|
||||
chunks = [c for c in chunks if c.strip() and len(c) > CHUNK_MIN_CHARS]
|
||||
return chunks
|
||||
|
||||
|
||||
def overlap_score(hypothesis_chunks, reference_chunks):
|
||||
if len(reference_chunks) > 0:
|
||||
length_modifier = len(hypothesis_chunks) / len(reference_chunks)
|
||||
else:
|
||||
length_modifier = 0
|
||||
search_distance = max(len(reference_chunks) // 5, 10)
|
||||
chunk_scores = []
|
||||
for i, hyp_chunk in enumerate(hypothesis_chunks):
|
||||
max_score = 0
|
||||
total_len = 0
|
||||
i_offset = int(i * length_modifier)
|
||||
chunk_range = range(max(0, i_offset-search_distance), min(len(reference_chunks), i_offset+search_distance))
|
||||
for j in chunk_range:
|
||||
ref_chunk = reference_chunks[j]
|
||||
score = fuzz.ratio(hyp_chunk, ref_chunk, score_cutoff=30) / 100
|
||||
if score > max_score:
|
||||
max_score = score
|
||||
total_len = len(ref_chunk)
|
||||
chunk_scores.append(max_score)
|
||||
return chunk_scores
|
||||
|
||||
|
||||
def score_text(hypothesis, reference):
|
||||
# Returns a 0-1 alignment score
|
||||
hypothesis_chunks = chunk_text(hypothesis)
|
||||
reference_chunks = chunk_text(reference)
|
||||
chunk_scores = overlap_score(hypothesis_chunks, reference_chunks)
|
||||
if len(chunk_scores) > 0:
|
||||
mean_score = mean(chunk_scores)
|
||||
return mean_score
|
||||
else:
|
||||
return 0
|
||||
#return mean(chunk_scores)
|
||||
@@ -1,9 +0,0 @@
|
||||
{
|
||||
"bucket_info":{
|
||||
"bucket-name-1":["ak", "sk", "endpoint"],
|
||||
"bucket-name-2":["ak", "sk", "endpoint"]
|
||||
},
|
||||
"temp-output-dir":"/tmp",
|
||||
"models-dir":"/tmp/models",
|
||||
"device-mode":"cpu"
|
||||
}
|
||||
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 541 KiB |
File diff suppressed because one or more lines are too long
Binary file not shown.
Binary file not shown.
@@ -1 +0,0 @@
|
||||
{"average_sim_score":0.6505598645664856, "average_edit_distance":0.2514908429188901, "average_bleu_score": 0.5808819533975296}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,54 +0,0 @@
|
||||
"""
|
||||
bench
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
from lib import calculate_score
|
||||
import pytest
|
||||
from conf import conf
|
||||
|
||||
code_path = os.environ.get('GITHUB_WORKSPACE')
|
||||
pdf_dev_path = conf.conf["pdf_dev_path"]
|
||||
pdf_res_path = conf.conf["pdf_res_path"]
|
||||
|
||||
class TestBench():
|
||||
"""
|
||||
test bench
|
||||
"""
|
||||
def test_ci_ben(self):
|
||||
"""
|
||||
ci benchmark
|
||||
"""
|
||||
fr = open(os.path.join(pdf_dev_path, "result.json"), "r", encoding="utf-8")
|
||||
lines = fr.readlines()
|
||||
last_line = lines[-1].strip()
|
||||
last_score = json.loads(last_line)
|
||||
last_simscore = last_score["average_sim_score"]
|
||||
last_editdistance = last_score["average_edit_distance"]
|
||||
last_bleu = last_score["average_bleu_score"]
|
||||
os.system(f"python tests/test_cli/lib/pre_clean.py --tool_name mineru --download_dir {pdf_dev_path}")
|
||||
now_score = get_score()
|
||||
print ("now_score:", now_score)
|
||||
if not os.path.exists(os.path.join(pdf_dev_path, "ci")):
|
||||
os.makedirs(os.path.join(pdf_dev_path, "ci"), exist_ok=True)
|
||||
fw = open(os.path.join(pdf_dev_path, "ci", "result.json"), "w+", encoding="utf-8")
|
||||
fw.write(json.dumps(now_score) + "\n")
|
||||
now_simscore = now_score["average_sim_score"]
|
||||
now_editdistance = now_score["average_edit_distance"]
|
||||
now_bleu = now_score["average_bleu_score"]
|
||||
assert last_simscore <= now_simscore
|
||||
assert last_editdistance <= now_editdistance
|
||||
assert last_bleu <= now_bleu
|
||||
|
||||
|
||||
def get_score():
|
||||
"""
|
||||
get score
|
||||
"""
|
||||
score = calculate_score.Scoring(os.path.join(pdf_dev_path, "result.json"))
|
||||
score.calculate_similarity_total("mineru", pdf_dev_path)
|
||||
res = score.summary_scores()
|
||||
return res
|
||||
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from conf import conf
|
||||
from lib import calculate_score
|
||||
|
||||
pdf_res_path = conf.conf['pdf_res_path']
|
||||
code_path = conf.conf['code_path']
|
||||
pdf_dev_path = conf.conf['pdf_dev_path']
|
||||
class TestCliCuda:
|
||||
"""test cli cuda."""
|
||||
def test_pdf_sdk_cuda(self):
|
||||
"""pdf sdk cuda."""
|
||||
clean_magicpdf(pdf_res_path)
|
||||
pdf_to_markdown()
|
||||
fr = open(os.path.join(pdf_dev_path, 'result.json'), 'r', encoding='utf-8')
|
||||
lines = fr.readlines()
|
||||
last_line = lines[-1].strip()
|
||||
last_score = json.loads(last_line)
|
||||
last_simscore = last_score['average_sim_score']
|
||||
last_editdistance = last_score['average_edit_distance']
|
||||
last_bleu = last_score['average_bleu_score']
|
||||
os.system(f'python tests/test_cli/lib/pre_clean.py --tool_name mineru --download_dir {pdf_dev_path}')
|
||||
now_score = get_score()
|
||||
print ('now_score:', now_score)
|
||||
if not os.path.exists(os.path.join(pdf_dev_path, 'ci')):
|
||||
os.makedirs(os.path.join(pdf_dev_path, 'ci'), exist_ok=True)
|
||||
fw = open(os.path.join(pdf_dev_path, 'ci', 'result.json'), 'w+', encoding='utf-8')
|
||||
fw.write(json.dumps(now_score) + '\n')
|
||||
now_simscore = now_score['average_sim_score']
|
||||
now_editdistance = now_score['average_edit_distance']
|
||||
now_bleu = now_score['average_bleu_score']
|
||||
assert last_simscore <= now_simscore
|
||||
assert last_editdistance <= now_editdistance
|
||||
assert last_bleu <= now_bleu
|
||||
|
||||
def pdf_to_markdown():
|
||||
"""pdf to md."""
|
||||
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:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'pdf', f'{demo_name}.pdf')
|
||||
cmd = 'magic-pdf pdf-command --pdf %s --inside_model true' % (pdf_path)
|
||||
os.system(cmd)
|
||||
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')
|
||||
src_path = os.path.join(pdf_res_path, demo_name, 'auto', f'{demo_name}.md')
|
||||
shutil.copy(src_path, res_path)
|
||||
|
||||
|
||||
|
||||
def get_score():
|
||||
"""get score."""
|
||||
score = calculate_score.Scoring(os.path.join(pdf_dev_path, 'result.json'))
|
||||
score.calculate_similarity_total('mineru', pdf_dev_path)
|
||||
res = score.summary_scores()
|
||||
return res
|
||||
|
||||
|
||||
def clean_magicpdf(pdf_res_path):
|
||||
"""clean magicpdf."""
|
||||
cmd = 'rm -rf %s' % (pdf_res_path)
|
||||
os.system(cmd)
|
||||
@@ -1,424 +0,0 @@
|
||||
"""test cli and sdk."""
|
||||
import logging
|
||||
import os
|
||||
import pytest
|
||||
from conf import conf
|
||||
from lib import common
|
||||
import time
|
||||
import magic_pdf.model as model_config
|
||||
from magic_pdf.data.read_api import read_local_images
|
||||
from magic_pdf.data.read_api import read_local_office
|
||||
from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter
|
||||
from magic_pdf.config.make_content_config import DropMode, MakeMode
|
||||
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
|
||||
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_local_sdk(self):
|
||||
"""pdf sdk auto test."""
|
||||
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:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'pdf', f'{demo_name}.pdf')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
name_without_suff = os.path.basename(pdf_path).split(".pdf")[0]
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
|
||||
reader1 = FileBasedDataReader("")
|
||||
pdf_bytes = reader1.read(pdf_path)
|
||||
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)
|
||||
common.delete_file(dir_path)
|
||||
### draw model result on each page
|
||||
infer_result.draw_model(os.path.join(dir_path, f"{name_without_suff}_model.pdf"))
|
||||
|
||||
### get model inference result
|
||||
model_inference_result = infer_result.get_infer_res()
|
||||
|
||||
### draw layout result on each page
|
||||
pipe_result.draw_layout(os.path.join(dir_path, f"{name_without_suff}_layout.pdf"))
|
||||
|
||||
### draw spans result on each page
|
||||
pipe_result.draw_span(os.path.join(dir_path, f"{name_without_suff}_spans.pdf"))
|
||||
|
||||
### dump markdown
|
||||
md_content = pipe_result.get_markdown(image_dir)
|
||||
pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir)
|
||||
### get content list content
|
||||
content_list_content = pipe_result.get_content_list(image_dir)
|
||||
pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir)
|
||||
|
||||
### get middle json
|
||||
middle_json_content = pipe_result.get_middle_json()
|
||||
### dump middle json
|
||||
pipe_result.dump_middle_json(md_writer, f'{name_without_suff}_middle.json')
|
||||
common.sdk_count_folders_and_check_contents(dir_path)
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_s3_sdk(self):
|
||||
"""pdf s3 sdk test."""
|
||||
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:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'pdf', f'{demo_name}.pdf')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
name_without_suff = os.path.basename(pdf_path).split(".pdf")[0]
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
pass
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_local_ppt(self):
|
||||
"""pdf sdk auto test."""
|
||||
demo_names = list()
|
||||
pdf_path = os.path.join(pdf_dev_path, 'ppt')
|
||||
for pdf_file in os.listdir(pdf_path):
|
||||
if pdf_file.endswith('.pptx'):
|
||||
demo_names.append(pdf_file.split('.')[0])
|
||||
for demo_name in demo_names:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'ppt', f'{demo_name}.pptx')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
name_without_suff = os.path.basename(pdf_path).split(".pptx")[0]
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
|
||||
ds = read_local_office(pdf_path)[0]
|
||||
common.delete_file(dir_path)
|
||||
|
||||
ds.apply(doc_analyze, ocr=True).pipe_txt_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir)
|
||||
common.sdk_count_folders_and_check_contents(dir_path)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_local_image(self):
|
||||
"""pdf sdk auto test."""
|
||||
demo_names = list()
|
||||
pdf_path = os.path.join(pdf_dev_path, 'images')
|
||||
for pdf_file in os.listdir(pdf_path):
|
||||
if pdf_file.endswith('.jpg'):
|
||||
demo_names.append(pdf_file.split('.')[0])
|
||||
for demo_name in demo_names:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'images', f'{demo_name}.jpg')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
name_without_suff = os.path.basename(pdf_path).split(".jpg")[0]
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
common.delete_file(dir_path)
|
||||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
|
||||
ds = read_local_images(pdf_path)[0]
|
||||
ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(
|
||||
md_writer, f"{name_without_suff}.md", image_dir)
|
||||
common.sdk_count_folders_and_check_contents(dir_path)
|
||||
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_local_image_dir(self):
|
||||
"""local image dir."""
|
||||
demo_names = list()
|
||||
pdf_path = os.path.join(pdf_dev_path, 'images')
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
|
||||
common.delete_file(dir_path)
|
||||
dss = read_local_images(pdf_path, suffixes=['.png', '.jpg'])
|
||||
count = 0
|
||||
for ds in dss:
|
||||
ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{count}.md", image_dir)
|
||||
count += 1
|
||||
common.sdk_count_folders_and_check_contents(dir_path)
|
||||
|
||||
def test_local_doc_parse(self):
|
||||
"""
|
||||
doc 解析
|
||||
"""
|
||||
demo_names = list()
|
||||
pdf_path = os.path.join(pdf_dev_path, 'doc')
|
||||
for pdf_file in os.listdir(pdf_path):
|
||||
if pdf_file.endswith('.docx'):
|
||||
demo_names.append(pdf_file.split('.')[0])
|
||||
for demo_name in demo_names:
|
||||
pdf_path = os.path.join(pdf_dev_path, 'doc', f'{demo_name}.docx')
|
||||
local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
name_without_suff = os.path.basename(pdf_path).split(".docx")[0]
|
||||
dir_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
|
||||
ds = read_local_office(pdf_path)[0]
|
||||
common.delete_file(dir_path)
|
||||
|
||||
ds.apply(doc_analyze, ocr=True).pipe_txt_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir)
|
||||
common.sdk_count_folders_and_check_contents(dir_path)
|
||||
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_cli_auto(self):
|
||||
"""magic_pdf cli test auto."""
|
||||
time.sleep(2)
|
||||
demo_names = []
|
||||
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:
|
||||
res_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
common.delete_file(res_path)
|
||||
cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
|
||||
pdf_path, f'{demo_name}.pdf'), res_path, 'auto')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
common.cli_count_folders_and_check_contents(
|
||||
os.path.join(res_path, demo_name, 'auto'))
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_cli_txt(self):
|
||||
"""magic_pdf cli test txt."""
|
||||
time.sleep(2)
|
||||
demo_names = []
|
||||
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:
|
||||
res_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
common.delete_file(res_path)
|
||||
cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
|
||||
pdf_path, f'{demo_name}.pdf'), res_path, 'txt')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
common.cli_count_folders_and_check_contents(
|
||||
os.path.join(res_path, demo_name, 'txt'))
|
||||
|
||||
@pytest.mark.P0
|
||||
def test_pdf_cli_ocr(self):
|
||||
"""magic_pdf cli test ocr."""
|
||||
time.sleep(2)
|
||||
demo_names = []
|
||||
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:
|
||||
res_path = os.path.join(pdf_dev_path, 'mineru')
|
||||
common.delete_file(res_path)
|
||||
cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
|
||||
pdf_path, f'{demo_name}.pdf'), res_path, 'ocr')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
common.cli_count_folders_and_check_contents(
|
||||
os.path.join(res_path, demo_name, 'ocr'))
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_local_jsonl_txt(self):
|
||||
"""magic_pdf_dev cli local txt."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, "txt")
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_local_jsonl_ocr(self):
|
||||
"""magic_pdf_dev cli local ocr."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'ocr')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_local_jsonl_auto(self):
|
||||
"""magic_pdf_dev cli local auto."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'auto')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_s3_jsonl_txt(self):
|
||||
"""magic_pdf_dev cli s3 txt."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, "txt")
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_s3_jsonl_ocr(self):
|
||||
"""magic_pdf_dev cli s3 ocr."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'ocr')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_s3_jsonl_auto(self):
|
||||
"""magic_pdf_dev cli s3 auto."""
|
||||
time.sleep(2)
|
||||
jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
|
||||
cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'auto')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_pdf_json_auto(self):
|
||||
"""magic_pdf_dev cli pdf+json auto."""
|
||||
time.sleep(2)
|
||||
json_path = os.path.join(pdf_dev_path, 'test_model.json')
|
||||
pdf_path = os.path.join(pdf_dev_path, 'pdf', 'test_rearch_report.pdf')
|
||||
cmd = 'magic-pdf-dev --pdf %s --json %s --method %s' % (pdf_path, json_path, 'auto')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.skip(reason='out-of-date api')
|
||||
@pytest.mark.P1
|
||||
def test_pdf_dev_cli_pdf_json_ocr(self):
|
||||
"""magic_pdf_dev cli pdf+json ocr."""
|
||||
time.sleep(2)
|
||||
json_path = os.path.join(pdf_dev_path, 'test_model.json')
|
||||
pdf_path = os.path.join(pdf_dev_path, 'pdf', 'test_rearch_report.pdf')
|
||||
cmd = 'magic-pdf-dev --pdf %s --json %s --method %s' % (pdf_path, json_path, 'auto')
|
||||
logging.info(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
@pytest.mark.P1
|
||||
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,
|
||||
"sub_model": "slanet_plus",
|
||||
"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.skip(reason="layoutlmv3废弃")
|
||||
@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_cpu(self):
|
||||
"""magic pdf cli cpu mode."""
|
||||
time.sleep(2)
|
||||
#pre_cmd = "cp ~/magic_pdf_html_table_cpu.json ~/magic-pdf.json"
|
||||
#os.system(pre_cmd)
|
||||
value = {
|
||||
"model": "rapid_table",
|
||||
"enable": True,
|
||||
"sub_model": "slanet_plus",
|
||||
"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)
|
||||
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_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)
|
||||
value = {
|
||||
"model": "rapid_table",
|
||||
"enable": False,
|
||||
"sub_model": "slanet_plus",
|
||||
"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_close_tables(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
|
||||
assert res is True
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main()
|
||||
|
||||
Binary file not shown.
@@ -1 +0,0 @@
|
||||
{"track_id":"e8824f5a-9fcb-4ee5-b2d4-6bf2c67019dc","path":"s3://sci-hub/enbook-scimag/78800000/libgen.scimag78872000-78872999/10.1017/cbo9780511770425.012.pdf","file_type":"pdf","content_type":"application/pdf","content_length":80078,"title":"German Idealism and the Concept of Punishment || Conclusion","remark":{"file_id":"scihub_78800000/libgen.scimag78872000-78872999.zip_10.1017/cbo9780511770425.012","file_source_type":"paper","original_file_id":"10.1017/cbo9780511770425.012","file_name":"10.1017/cbo9780511770425.012.pdf","author":"Merle, Jean-Christophe"}}
|
||||
@@ -1 +0,0 @@
|
||||
{"track_id":"e8824f5a-9fcb-4ee5-b2d4-6bf2c67019dc","path":"tests/unittest/test_data/assets/pdfs/test_02.pdf","file_type":"pdf","content_type":"application/pdf","content_length":80078,"title":"German Idealism and the Concept of Punishment || Conclusion","remark":{"file_id":"scihub_78800000/libgen.scimag78872000-78872999.zip_10.1017/cbo9780511770425.012","file_source_type":"paper","original_file_id":"10.1017/cbo9780511770425.012","file_name":"10.1017/cbo9780511770425.012.pdf","author":"Merle, Jean-Christophe"}}
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 326 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 410 KiB |
@@ -1,24 +0,0 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from magic_pdf.data.data_reader_writer import (FileBasedDataReader,
|
||||
FileBasedDataWriter)
|
||||
|
||||
|
||||
def test_filebased_reader_writer():
|
||||
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/data/filebased_reader_writer'
|
||||
sub_dir = os.path.join(unitest_dir, 'sub')
|
||||
abs_fn = os.path.join(unitest_dir, 'abspath.txt')
|
||||
|
||||
os.makedirs(sub_dir, exist_ok=True)
|
||||
|
||||
writer = FileBasedDataWriter(sub_dir)
|
||||
reader = FileBasedDataReader(sub_dir)
|
||||
|
||||
writer.write('test.txt', b'hello world')
|
||||
assert reader.read('test.txt') == b'hello world'
|
||||
|
||||
writer.write(abs_fn, b'hello world')
|
||||
assert reader.read(abs_fn) == b'hello world'
|
||||
shutil.rmtree(unitest_dir)
|
||||
@@ -1,160 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import fitz
|
||||
import pytest
|
||||
|
||||
from magic_pdf.data.data_reader_writer import (MultiBucketS3DataReader,
|
||||
MultiBucketS3DataWriter)
|
||||
from magic_pdf.data.schemas import S3Config
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY_2', None) is None, reason='need s3 config!'
|
||||
)
|
||||
def test_multi_bucket_s3_reader_writer():
|
||||
"""test multi bucket s3 reader writer must config s3 config in the
|
||||
environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
|
||||
S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx.
|
||||
|
||||
export S3_BUCKET_2=xxx export S3_ACCESS_KEY_2=xxx export S3_SECRET_KEY_2=xxx export S3_ENDPOINT_2=xxx
|
||||
"""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
|
||||
bucket_2 = os.getenv('S3_BUCKET_2', '')
|
||||
ak_2 = os.getenv('S3_ACCESS_KEY_2', '')
|
||||
sk_2 = os.getenv('S3_SECRET_KEY_2', '')
|
||||
endpoint_url_2 = os.getenv('S3_ENDPOINT_2', '')
|
||||
|
||||
s3configs = [
|
||||
S3Config(
|
||||
bucket_name=bucket, access_key=ak, secret_key=sk, endpoint_url=endpoint_url
|
||||
),
|
||||
S3Config(
|
||||
bucket_name=bucket_2,
|
||||
access_key=ak_2,
|
||||
secret_key=sk_2,
|
||||
endpoint_url=endpoint_url_2,
|
||||
),
|
||||
]
|
||||
|
||||
reader = MultiBucketS3DataReader(bucket, s3configs)
|
||||
writer = MultiBucketS3DataWriter(bucket, s3configs)
|
||||
|
||||
bits = reader.read('meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl')
|
||||
|
||||
assert bits == reader.read(
|
||||
f's3://{bucket}/meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl'
|
||||
)
|
||||
|
||||
bits = reader.read(
|
||||
f's3://{bucket_2}/enbook-scimag/78800000/libgen.scimag78872000-78872999/10.1017/cbo9780511770425.012.pdf'
|
||||
)
|
||||
docs = fitz.open('pdf', bits)
|
||||
assert len(docs) == 10
|
||||
|
||||
bits = reader.read(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl?bytes=566,713'
|
||||
)
|
||||
assert bits == reader.read_at(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl', 566, 713
|
||||
)
|
||||
assert len(json.loads(bits)) > 0
|
||||
|
||||
writer.write_string(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt', 'abc'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
writer.write(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt',
|
||||
'123'.encode(),
|
||||
)
|
||||
|
||||
assert '123'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY_2', None) is None, reason='need s3 config!'
|
||||
)
|
||||
def test_multi_bucket_s3_reader_writer_with_prefix():
|
||||
"""test multi bucket s3 reader writer must config s3 config in the
|
||||
environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
|
||||
S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx.
|
||||
|
||||
export S3_BUCKET_2=xxx export S3_ACCESS_KEY_2=xxx export S3_SECRET_KEY_2=xxx export S3_ENDPOINT_2=xxx
|
||||
"""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
|
||||
bucket_2 = os.getenv('S3_BUCKET_2', '')
|
||||
ak_2 = os.getenv('S3_ACCESS_KEY_2', '')
|
||||
sk_2 = os.getenv('S3_SECRET_KEY_2', '')
|
||||
endpoint_url_2 = os.getenv('S3_ENDPOINT_2', '')
|
||||
|
||||
s3configs = [
|
||||
S3Config(
|
||||
bucket_name=bucket, access_key=ak, secret_key=sk, endpoint_url=endpoint_url
|
||||
),
|
||||
S3Config(
|
||||
bucket_name=bucket_2,
|
||||
access_key=ak_2,
|
||||
secret_key=sk_2,
|
||||
endpoint_url=endpoint_url_2,
|
||||
),
|
||||
]
|
||||
|
||||
prefix = 'meta-index'
|
||||
reader = MultiBucketS3DataReader(f'{bucket}/{prefix}', s3configs)
|
||||
writer = MultiBucketS3DataWriter(f'{bucket}/{prefix}', s3configs)
|
||||
|
||||
bits = reader.read('scihub/v001/scihub/part-66210c190659-000026.jsonl')
|
||||
|
||||
assert bits == reader.read(
|
||||
f's3://{bucket}/{prefix}/scihub/v001/scihub/part-66210c190659-000026.jsonl'
|
||||
)
|
||||
|
||||
bits = reader.read(
|
||||
f's3://{bucket_2}/enbook-scimag/78800000/libgen.scimag78872000-78872999/10.1017/cbo9780511770425.012.pdf'
|
||||
)
|
||||
docs = fitz.open('pdf', bits)
|
||||
assert len(docs) == 10
|
||||
|
||||
bits = reader.read(
|
||||
'scihub/v001/scihub/part-66210c190659-000026.jsonl?bytes=566,713'
|
||||
)
|
||||
assert bits == reader.read_at(
|
||||
'scihub/v001/scihub/part-66210c190659-000026.jsonl', 566, 713
|
||||
)
|
||||
assert len(json.loads(bits)) > 0
|
||||
|
||||
writer.write_string(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt', 'abc'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
f's3://{bucket}/{prefix}/unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
writer.write(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt',
|
||||
'123'.encode(),
|
||||
)
|
||||
|
||||
assert '123'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt'
|
||||
)
|
||||
@@ -1,106 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY', None) is None, reason='need s3 config!'
|
||||
)
|
||||
def test_s3_reader_writer():
|
||||
"""test multi bucket s3 reader writer must config s3 config in the
|
||||
environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
|
||||
S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx."""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
|
||||
reader = S3DataReader('', bucket, ak, sk, endpoint_url)
|
||||
writer = S3DataWriter('', bucket, ak, sk, endpoint_url)
|
||||
|
||||
bits = reader.read('meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl')
|
||||
|
||||
assert bits == reader.read(
|
||||
f's3://{bucket}/meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl'
|
||||
)
|
||||
|
||||
bits = reader.read(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl?bytes=566,713'
|
||||
)
|
||||
assert bits == reader.read_at(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl', 566, 713
|
||||
)
|
||||
assert len(json.loads(bits)) > 0
|
||||
|
||||
writer.write_string(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt', 'abc'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
writer.write(
|
||||
f'{bucket}/unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt',
|
||||
'123'.encode(),
|
||||
)
|
||||
|
||||
assert '123'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY', None) is None, reason='need s3 config!'
|
||||
)
|
||||
def test_s3_reader_writer_with_prefix():
|
||||
"""test multi bucket s3 reader writer must config s3 config in the
|
||||
environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
|
||||
S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx."""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
|
||||
prefix = 'meta-index'
|
||||
|
||||
reader = S3DataReader(prefix, bucket, ak, sk, endpoint_url)
|
||||
writer = S3DataWriter(prefix, bucket, ak, sk, endpoint_url)
|
||||
|
||||
bits = reader.read('scihub/v001/scihub/part-66210c190659-000026.jsonl')
|
||||
|
||||
assert bits == reader.read(
|
||||
f's3://{bucket}/{prefix}/scihub/v001/scihub/part-66210c190659-000026.jsonl'
|
||||
)
|
||||
|
||||
bits = reader.read(
|
||||
'scihub/v001/scihub/part-66210c190659-000026.jsonl?bytes=566,713'
|
||||
)
|
||||
assert bits == reader.read_at(
|
||||
'scihub/v001/scihub/part-66210c190659-000026.jsonl', 566, 713
|
||||
)
|
||||
assert len(json.loads(bits)) > 0
|
||||
|
||||
writer.write_string(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt', 'abc'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
assert 'abc'.encode() == reader.read(
|
||||
f's3://{bucket}/{prefix}/unittest/data/data_reader_writer/multi_bucket_s3_data/test01.txt'
|
||||
)
|
||||
|
||||
writer.write(
|
||||
f'{bucket}/{prefix}/unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt',
|
||||
'123'.encode(),
|
||||
)
|
||||
|
||||
assert '123'.encode() == reader.read(
|
||||
'unittest/data/data_reader_writer/multi_bucket_s3_data/test02.txt'
|
||||
)
|
||||
@@ -1,55 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.data.io.s3 import S3Reader, S3Writer
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY', None) is None, reason='s3 config not found'
|
||||
)
|
||||
def test_s3_reader():
|
||||
"""test s3 reader.
|
||||
|
||||
must config s3 config in the environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export S3_SECRET_KEY=xxx
|
||||
export S3_ENDPOINT=xxx
|
||||
"""
|
||||
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
reader = S3Reader(bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint_url)
|
||||
bits = reader.read(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl'
|
||||
)
|
||||
assert len(bits) > 0
|
||||
|
||||
bits = reader.read_at(
|
||||
'meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl',
|
||||
566,
|
||||
713,
|
||||
)
|
||||
assert len(json.loads(bits)) > 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY', None) is None, reason='s3 config not found'
|
||||
)
|
||||
def test_s3_writer():
|
||||
"""test s3 reader.
|
||||
|
||||
must config s3 config in the environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export S3_SECRET_KEY=xxx
|
||||
export S3_ENDPOINT=xxx
|
||||
"""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
writer = S3Writer(bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint_url)
|
||||
test_fn = 'unittest/io/test.jsonl'
|
||||
writer.write(test_fn, '123'.encode())
|
||||
reader = S3Reader(bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint_url)
|
||||
bits = reader.read(test_fn)
|
||||
assert bits.decode() == '123'
|
||||
@@ -1,18 +0,0 @@
|
||||
|
||||
from magic_pdf.data.dataset import ImageDataset, PymuDocDataset
|
||||
|
||||
|
||||
def test_pymudataset():
|
||||
with open('tests/unittest/test_data/assets/pdfs/test_01.pdf', 'rb') as f:
|
||||
bits = f.read()
|
||||
datasets = PymuDocDataset(bits)
|
||||
assert len(datasets) > 0
|
||||
assert datasets.get_page(0).get_page_info().h > 100
|
||||
|
||||
|
||||
def test_imagedataset():
|
||||
with open('tests/unittest/test_data/assets/pngs/test_01.png', 'rb') as f:
|
||||
bits = f.read()
|
||||
datasets = ImageDataset(bits)
|
||||
assert len(datasets) == 1
|
||||
assert datasets.get_page(0).get_page_info().w > 100
|
||||
@@ -1,115 +0,0 @@
|
||||
import pytest
|
||||
import json
|
||||
from magic_pdf.libs.json_compressor import JsonCompressor
|
||||
|
||||
# Test data fixtures
|
||||
@pytest.fixture
|
||||
def test_cases():
|
||||
return [
|
||||
# Simple dictionary
|
||||
{"name": "John", "age": 30},
|
||||
|
||||
# Nested dictionary
|
||||
{
|
||||
"person": {
|
||||
"name": "Alice",
|
||||
"address": {
|
||||
"street": "123 Main St",
|
||||
"city": "New York"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
# List of dictionaries
|
||||
[
|
||||
{"id": 1, "value": "first"},
|
||||
{"id": 2, "value": "second"}
|
||||
],
|
||||
|
||||
# Dictionary with various data types
|
||||
{
|
||||
"string": "hello",
|
||||
"integer": 42,
|
||||
"float": 3.14,
|
||||
"boolean": True,
|
||||
"null": None,
|
||||
"array": [1, 2, 3],
|
||||
"nested": {"key": "value"}
|
||||
},
|
||||
|
||||
# Empty structures
|
||||
{},
|
||||
[],
|
||||
{"empty_list": [], "empty_dict": {}}
|
||||
]
|
||||
|
||||
@pytest.fixture
|
||||
def large_data():
|
||||
return {
|
||||
"data": ["test" * 100] * 100 # Create a large repeated string
|
||||
}
|
||||
|
||||
def test_compression_decompression_cycle(test_cases):
|
||||
"""Test that data remains intact after compression and decompression"""
|
||||
for test_data in test_cases:
|
||||
# Compress the data
|
||||
compressed = JsonCompressor.compress_json(test_data)
|
||||
|
||||
# Verify compressed string is not empty and is a string
|
||||
assert isinstance(compressed, str)
|
||||
assert len(compressed) > 0
|
||||
|
||||
# Decompress the data
|
||||
decompressed = JsonCompressor.decompress_json(compressed)
|
||||
|
||||
# Verify the decompressed data matches original
|
||||
assert test_data == decompressed
|
||||
|
||||
def test_compression_reduces_size(large_data):
|
||||
"""Test that compression actually reduces data size for large enough input"""
|
||||
original_size = len(json.dumps(large_data))
|
||||
compressed = JsonCompressor.compress_json(large_data)
|
||||
compressed_size = len(compressed)
|
||||
|
||||
# Verify compression actually saved space
|
||||
assert compressed_size < original_size
|
||||
|
||||
def test_invalid_json_serializable():
|
||||
"""Test handling of non-JSON serializable input"""
|
||||
with pytest.raises(TypeError):
|
||||
JsonCompressor.compress_json(set([1, 2, 3])) # sets are not JSON serializable
|
||||
|
||||
def test_invalid_compressed_string():
|
||||
"""Test handling of invalid compressed string"""
|
||||
with pytest.raises(Exception):
|
||||
JsonCompressor.decompress_json("invalid_base64_string")
|
||||
|
||||
def test_empty_string_input():
|
||||
"""Test handling of empty string input"""
|
||||
with pytest.raises(Exception):
|
||||
JsonCompressor.decompress_json("")
|
||||
|
||||
def test_special_characters():
|
||||
"""Test handling of special characters"""
|
||||
test_data = {
|
||||
"special": "!@#$%^&*()_+-=[]{}|;:,.<>?",
|
||||
"unicode": "Hello 世界 🌍"
|
||||
}
|
||||
|
||||
compressed = JsonCompressor.compress_json(test_data)
|
||||
decompressed = JsonCompressor.decompress_json(compressed)
|
||||
assert test_data == decompressed
|
||||
|
||||
# Parametrized test for different types of input
|
||||
@pytest.mark.parametrize("test_input", [
|
||||
{"simple": "value"},
|
||||
[1, 2, 3],
|
||||
{"nested": {"key": "value"}},
|
||||
["mixed", 1, True, None],
|
||||
{"unicode": "🌍"}
|
||||
])
|
||||
def test_various_input_types(test_input):
|
||||
"""Test compression and decompression with various input types"""
|
||||
compressed = JsonCompressor.compress_json(test_input)
|
||||
decompressed = JsonCompressor.decompress_json(compressed)
|
||||
assert test_input == decompressed
|
||||
@@ -1,78 +0,0 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.data.data_reader_writer import MultiBucketS3DataReader
|
||||
from magic_pdf.data.read_api import (read_jsonl, read_local_images,
|
||||
read_local_pdfs)
|
||||
from magic_pdf.data.schemas import S3Config
|
||||
|
||||
|
||||
def test_read_local_pdfs():
|
||||
datasets = read_local_pdfs('tests/unittest/test_data/assets/pdfs')
|
||||
assert len(datasets) == 2
|
||||
assert len(datasets[0]) > 0
|
||||
assert len(datasets[1]) > 0
|
||||
|
||||
assert datasets[0].get_page(0).get_page_info().w > 0
|
||||
assert datasets[0].get_page(0).get_page_info().h > 0
|
||||
|
||||
|
||||
def test_read_local_images():
|
||||
datasets = read_local_images('tests/unittest/test_data/assets/pngs', suffixes=['.png'])
|
||||
assert len(datasets) == 2
|
||||
assert len(datasets[0]) == 1
|
||||
assert len(datasets[1]) == 1
|
||||
|
||||
assert datasets[0].get_page(0).get_page_info().w > 0
|
||||
assert datasets[0].get_page(0).get_page_info().h > 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.getenv('S3_ACCESS_KEY_2', None) is None, reason='need s3 config!'
|
||||
)
|
||||
def test_read_json():
|
||||
"""test multi bucket s3 reader writer must config s3 config in the
|
||||
environment export S3_BUCKET=xxx export S3_ACCESS_KEY=xxx export
|
||||
S3_SECRET_KEY=xxx export S3_ENDPOINT=xxx.
|
||||
|
||||
export S3_BUCKET_2=xxx export S3_ACCESS_KEY_2=xxx export S3_SECRET_KEY_2=xxx export S3_ENDPOINT_2=xxx
|
||||
"""
|
||||
bucket = os.getenv('S3_BUCKET', '')
|
||||
ak = os.getenv('S3_ACCESS_KEY', '')
|
||||
sk = os.getenv('S3_SECRET_KEY', '')
|
||||
endpoint_url = os.getenv('S3_ENDPOINT', '')
|
||||
|
||||
bucket_2 = os.getenv('S3_BUCKET_2', '')
|
||||
ak_2 = os.getenv('S3_ACCESS_KEY_2', '')
|
||||
sk_2 = os.getenv('S3_SECRET_KEY_2', '')
|
||||
endpoint_url_2 = os.getenv('S3_ENDPOINT_2', '')
|
||||
|
||||
s3configs = [
|
||||
S3Config(
|
||||
bucket_name=bucket, access_key=ak, secret_key=sk, endpoint_url=endpoint_url
|
||||
),
|
||||
S3Config(
|
||||
bucket_name=bucket_2,
|
||||
access_key=ak_2,
|
||||
secret_key=sk_2,
|
||||
endpoint_url=endpoint_url_2,
|
||||
),
|
||||
]
|
||||
|
||||
reader = MultiBucketS3DataReader(bucket, s3configs)
|
||||
|
||||
datasets = read_jsonl(
|
||||
f's3://{bucket}/meta-index/scihub/v001/scihub/part-66210c190659-000026.jsonl',
|
||||
reader,
|
||||
)
|
||||
assert len(datasets) > 0
|
||||
assert len(datasets[0]) == 10
|
||||
|
||||
datasets = read_jsonl('tests/unittest/test_data/assets/jsonl/test_01.jsonl', reader)
|
||||
assert len(datasets) == 1
|
||||
assert len(datasets[0]) == 10
|
||||
|
||||
datasets = read_jsonl('tests/unittest/test_data/assets/jsonl/test_02.jsonl')
|
||||
assert len(datasets) == 1
|
||||
assert len(datasets[0]) == 1
|
||||
@@ -0,0 +1,293 @@
|
||||
# Copyright (c) Opendatalab. All rights reserved.
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from bs4 import BeautifulSoup
|
||||
from fuzzywuzzy import fuzz
|
||||
from mineru.cli.common import (
|
||||
convert_pdf_bytes_to_bytes_by_pypdfium2,
|
||||
prepare_env,
|
||||
read_fn,
|
||||
)
|
||||
from mineru.data.data_reader_writer import FileBasedDataWriter
|
||||
from mineru.utils.enum_class import MakeMode
|
||||
from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
|
||||
from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
|
||||
from mineru.backend.pipeline.pipeline_middle_json_mkcontent import (
|
||||
union_make as pipeline_union_make,
|
||||
)
|
||||
from mineru.backend.pipeline.model_json_to_middle_json import (
|
||||
result_to_middle_json as pipeline_result_to_middle_json,
|
||||
)
|
||||
from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
|
||||
|
||||
|
||||
def test_pipeline_with_two_config():
|
||||
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||
pdf_files_dir = os.path.join(__dir__, "pdfs")
|
||||
output_dir = os.path.join(__dir__, "output")
|
||||
pdf_suffixes = [".pdf"]
|
||||
image_suffixes = [".png", ".jpeg", ".jpg"]
|
||||
|
||||
doc_path_list = []
|
||||
for doc_path in Path(pdf_files_dir).glob("*"):
|
||||
if doc_path.suffix in pdf_suffixes + image_suffixes:
|
||||
doc_path_list.append(doc_path)
|
||||
|
||||
os.environ["MINERU_MODEL_SOURCE"] = "modelscope"
|
||||
|
||||
pdf_file_names = []
|
||||
pdf_bytes_list = []
|
||||
p_lang_list = []
|
||||
for path in doc_path_list:
|
||||
file_name = str(Path(path).stem)
|
||||
pdf_bytes = read_fn(path)
|
||||
pdf_file_names.append(file_name)
|
||||
pdf_bytes_list.append(pdf_bytes)
|
||||
p_lang_list.append("en")
|
||||
for idx, pdf_bytes in enumerate(pdf_bytes_list):
|
||||
new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes)
|
||||
pdf_bytes_list[idx] = new_pdf_bytes
|
||||
|
||||
# 获取 pipline 分析结果, 分别测试 txt 和 ocr 两种解析方法的结果
|
||||
infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = (
|
||||
pipeline_doc_analyze(
|
||||
pdf_bytes_list,
|
||||
p_lang_list,
|
||||
parse_method="txt",
|
||||
)
|
||||
)
|
||||
write_infer_result(
|
||||
infer_results,
|
||||
all_image_lists,
|
||||
all_pdf_docs,
|
||||
lang_list,
|
||||
ocr_enabled_list,
|
||||
pdf_file_names,
|
||||
output_dir,
|
||||
parse_method="txt",
|
||||
)
|
||||
assert_content("tests/unittest/output/test/txt/test_content_list.json")
|
||||
infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = (
|
||||
pipeline_doc_analyze(
|
||||
pdf_bytes_list,
|
||||
p_lang_list,
|
||||
parse_method="ocr",
|
||||
)
|
||||
)
|
||||
write_infer_result(
|
||||
infer_results,
|
||||
all_image_lists,
|
||||
all_pdf_docs,
|
||||
lang_list,
|
||||
ocr_enabled_list,
|
||||
pdf_file_names,
|
||||
output_dir,
|
||||
parse_method="ocr",
|
||||
)
|
||||
assert_content("tests/unittest/output/test/ocr/test_content_list.json")
|
||||
|
||||
|
||||
def test_vlm_transformers_with_default_config():
|
||||
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||
pdf_files_dir = os.path.join(__dir__, "pdfs")
|
||||
output_dir = os.path.join(__dir__, "output")
|
||||
pdf_suffixes = [".pdf"]
|
||||
image_suffixes = [".png", ".jpeg", ".jpg"]
|
||||
|
||||
doc_path_list = []
|
||||
for doc_path in Path(pdf_files_dir).glob("*"):
|
||||
if doc_path.suffix in pdf_suffixes + image_suffixes:
|
||||
doc_path_list.append(doc_path)
|
||||
|
||||
os.environ["MINERU_MODEL_SOURCE"] = "modelscope"
|
||||
|
||||
pdf_file_names = []
|
||||
pdf_bytes_list = []
|
||||
p_lang_list = []
|
||||
for path in doc_path_list:
|
||||
file_name = str(Path(path).stem)
|
||||
pdf_bytes = read_fn(path)
|
||||
pdf_file_names.append(file_name)
|
||||
pdf_bytes_list.append(pdf_bytes)
|
||||
p_lang_list.append("en")
|
||||
|
||||
for idx, pdf_bytes in enumerate(pdf_bytes_list):
|
||||
pdf_file_name = pdf_file_names[idx]
|
||||
pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes)
|
||||
local_image_dir, local_md_dir = prepare_env(
|
||||
output_dir, pdf_file_name, parse_method="vlm"
|
||||
)
|
||||
image_writer, md_writer = FileBasedDataWriter(
|
||||
local_image_dir
|
||||
), FileBasedDataWriter(local_md_dir)
|
||||
middle_json, infer_result = vlm_doc_analyze(
|
||||
pdf_bytes, image_writer=image_writer, backend="transformers"
|
||||
)
|
||||
|
||||
pdf_info = middle_json["pdf_info"]
|
||||
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
|
||||
md_content_str = vlm_union_make(pdf_info, MakeMode.MM_MD, image_dir)
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}.md",
|
||||
md_content_str,
|
||||
)
|
||||
|
||||
content_list = vlm_union_make(pdf_info, MakeMode.CONTENT_LIST, image_dir)
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_content_list.json",
|
||||
json.dumps(content_list, ensure_ascii=False, indent=4),
|
||||
)
|
||||
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_middle.json",
|
||||
json.dumps(middle_json, ensure_ascii=False, indent=4),
|
||||
)
|
||||
|
||||
model_output = ("\n" + "-" * 50 + "\n").join(infer_result)
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_model_output.txt",
|
||||
model_output,
|
||||
)
|
||||
|
||||
logger.info(f"local output dir is {local_md_dir}")
|
||||
assert_content("tests/unittest/output/test/vlm/test_content_list.json")
|
||||
|
||||
|
||||
def write_infer_result(
|
||||
infer_results,
|
||||
all_image_lists,
|
||||
all_pdf_docs,
|
||||
lang_list,
|
||||
ocr_enabled_list,
|
||||
pdf_file_names,
|
||||
output_dir,
|
||||
parse_method,
|
||||
):
|
||||
for idx, model_list in enumerate(infer_results):
|
||||
model_json = copy.deepcopy(model_list)
|
||||
pdf_file_name = pdf_file_names[idx]
|
||||
local_image_dir, local_md_dir = prepare_env(
|
||||
output_dir, pdf_file_name, parse_method
|
||||
)
|
||||
image_writer, md_writer = FileBasedDataWriter(
|
||||
local_image_dir
|
||||
), FileBasedDataWriter(local_md_dir)
|
||||
|
||||
images_list = all_image_lists[idx]
|
||||
pdf_doc = all_pdf_docs[idx]
|
||||
_lang = lang_list[idx]
|
||||
_ocr_enable = ocr_enabled_list[idx]
|
||||
middle_json = pipeline_result_to_middle_json(
|
||||
model_list,
|
||||
images_list,
|
||||
pdf_doc,
|
||||
image_writer,
|
||||
_lang,
|
||||
_ocr_enable,
|
||||
True,
|
||||
)
|
||||
|
||||
pdf_info = middle_json["pdf_info"]
|
||||
|
||||
image_dir = str(os.path.basename(local_image_dir))
|
||||
# 写入 md 文件
|
||||
md_content_str = pipeline_union_make(pdf_info, MakeMode.MM_MD, image_dir)
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}.md",
|
||||
md_content_str,
|
||||
)
|
||||
|
||||
content_list = pipeline_union_make(pdf_info, MakeMode.CONTENT_LIST, image_dir)
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_content_list.json",
|
||||
json.dumps(content_list, ensure_ascii=False, indent=4),
|
||||
)
|
||||
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_middle.json",
|
||||
json.dumps(middle_json, ensure_ascii=False, indent=4),
|
||||
)
|
||||
|
||||
md_writer.write_string(
|
||||
f"{pdf_file_name}_model.json",
|
||||
json.dumps(model_json, ensure_ascii=False, indent=4),
|
||||
)
|
||||
|
||||
logger.info(f"local output dir is {local_md_dir}")
|
||||
|
||||
|
||||
def validate_html(html_content):
|
||||
try:
|
||||
soup = BeautifulSoup(html_content, "html.parser")
|
||||
return True
|
||||
except Exception as e:
|
||||
return False
|
||||
|
||||
|
||||
def assert_content(content_path):
|
||||
content_list = []
|
||||
with open(content_path, "r", encoding="utf-8") as file:
|
||||
content_list = json.load(file)
|
||||
type_set = set()
|
||||
for content_dict in content_list:
|
||||
match content_dict["type"]:
|
||||
# 图片校验，只校验 Caption
|
||||
case "image":
|
||||
type_set.add("image")
|
||||
assert (
|
||||
content_dict["image_caption"][0].strip().lower()
|
||||
== "Figure 1: Figure Caption".lower()
|
||||
)
|
||||
# 表格校验，校验 Caption，表格格式和表格内容
|
||||
case "table":
|
||||
type_set.add("table")
|
||||
assert (
|
||||
content_dict["table_caption"][0].strip().lower()
|
||||
== "Table 1: Table Caption".lower()
|
||||
)
|
||||
assert validate_html(content_dict["table_body"])
|
||||
target_str_list = [
|
||||
"Linear Regression",
|
||||
"0.98740",
|
||||
"1321.2",
|
||||
"2-order Polynomial",
|
||||
"0.99906",
|
||||
"26.4",
|
||||
"3-order Polynomial",
|
||||
"0.99913",
|
||||
"101.2",
|
||||
"4-order Polynomial",
|
||||
"0.99914",
|
||||
"94.1",
|
||||
"Gray Prediction",
|
||||
"0.00617",
|
||||
"687",
|
||||
]
|
||||
correct_count = 0
|
||||
for target_str in target_str_list:
|
||||
if target_str in content_dict["table_body"]:
|
||||
correct_count += 1
|
||||
|
||||
assert correct_count > 0.9 * len(target_str_list)
|
||||
# 公式校验，检测是否含有公式元素
|
||||
case "equation":
|
||||
type_set.add("equation")
|
||||
target_str_list = ["$$", "lambda", "frac", "bar"]
|
||||
for target_str in target_str_list:
|
||||
assert target_str in content_dict["text"]
|
||||
# 文本校验，文本相似度超过90
|
||||
case "text":
|
||||
type_set.add("text")
|
||||
assert (
|
||||
fuzz.ratio(
|
||||
content_dict["text"],
|
||||
"Trump graduated from the Wharton School of the University of Pennsylvania with a bachelor's degree in 1968. He became president of his father's real estate business in 1971 and renamed it The Trump Organization.",
|
||||
)
|
||||
> 90
|
||||
)
|
||||
assert len(type_set) >= 4
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
@@ -1,55 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from magic_pdf.integrations.rag.api import DataReader, RagDocumentReader
|
||||
from magic_pdf.integrations.rag.type import CategoryType
|
||||
from magic_pdf.integrations.rag.utils import \
|
||||
convert_middle_json_to_layout_elements
|
||||
|
||||
|
||||
def test_rag_document_reader():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/integrations/rag'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir=unitest_dir)
|
||||
os.makedirs(temp_output_dir, exist_ok=True)
|
||||
|
||||
# test
|
||||
with open('tests/unittest/test_integrations/test_rag/assets/middle.json') as f:
|
||||
json_data = json.load(f)
|
||||
res = convert_middle_json_to_layout_elements(json_data, temp_output_dir)
|
||||
|
||||
doc = RagDocumentReader(res)
|
||||
assert len(list(iter(doc))) == 1
|
||||
|
||||
page = list(iter(doc))[0]
|
||||
assert len(list(iter(page))) >= 10
|
||||
assert len(page.get_rel_map()) >= 3
|
||||
|
||||
item = list(iter(page))[0]
|
||||
assert item.category_type == CategoryType.text
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
|
||||
|
||||
def test_data_reader():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/integrations/rag'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir=unitest_dir)
|
||||
os.makedirs(temp_output_dir, exist_ok=True)
|
||||
|
||||
# test
|
||||
data_reader = DataReader('tests/unittest/test_integrations/test_rag/assets', 'ocr',
|
||||
temp_output_dir)
|
||||
|
||||
assert data_reader.get_documents_count() == 2
|
||||
for idx in range(data_reader.get_documents_count()):
|
||||
document = data_reader.get_document_result(idx)
|
||||
assert document is not None
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
@@ -1,57 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from magic_pdf.integrations.rag.type import CategoryType
|
||||
from magic_pdf.integrations.rag.utils import (
|
||||
convert_middle_json_to_layout_elements, inference)
|
||||
|
||||
|
||||
def test_convert_middle_json_to_layout_elements():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/integrations/rag'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir=unitest_dir)
|
||||
os.makedirs(temp_output_dir, exist_ok=True)
|
||||
|
||||
# test
|
||||
with open('tests/unittest/test_integrations/test_rag/assets/middle.json') as f:
|
||||
json_data = json.load(f)
|
||||
res = convert_middle_json_to_layout_elements(json_data, temp_output_dir)
|
||||
|
||||
assert len(res) == 1
|
||||
assert len(res[0].layout_dets) > 0
|
||||
assert res[0].layout_dets[0].anno_id == 0
|
||||
assert res[0].layout_dets[0].category_type == CategoryType.text
|
||||
assert len(res[0].extra.element_relation) >= 2
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
|
||||
|
||||
def test_inference():
|
||||
|
||||
asset_dir = 'tests/unittest/test_integrations/test_rag/assets'
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/integrations/rag'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir=unitest_dir)
|
||||
os.makedirs(temp_output_dir, exist_ok=True)
|
||||
|
||||
# test
|
||||
res = inference(
|
||||
asset_dir + '/one_page_with_table_image.pdf',
|
||||
temp_output_dir,
|
||||
'ocr',
|
||||
)
|
||||
|
||||
assert res is not None
|
||||
assert len(res) == 1
|
||||
assert len(res[0].layout_dets) > 0
|
||||
assert res[0].layout_dets[0].anno_id == 0
|
||||
assert res[0].layout_dets[0].category_type == CategoryType.text
|
||||
assert len(res[0].extra.element_relation) >= 2
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
@@ -1,140 +0,0 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.filter.pdf_classify_by_type import classify_by_area, classify_by_text_len, classify_by_avg_words, \
|
||||
classify_by_img_num, classify_by_text_layout, classify_by_img_narrow_strips
|
||||
from magic_pdf.filter.pdf_meta_scan import get_pdf_page_size_pts, get_pdf_textlen_per_page, get_imgs_per_page
|
||||
from test_commons import get_docs_from_test_pdf, get_test_json_data
|
||||
|
||||
# 获取当前目录
|
||||
current_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
'''
|
||||
根据图片尺寸占页面面积的比例，判断是否为扫描版
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_area",
|
||||
[
|
||||
("the_eye/the_eye_cdn_00391653", True), # 特殊文字版1.每页存储所有图片，特点是图片占页面比例不大，每页展示可能为0也可能不止1张
|
||||
("scihub/scihub_08400000/libgen.scimag08489000-08489999.zip_10.1016/0370-1573(90)90070-i", False), # 特殊扫描版2，每页存储的扫描页图片数量递增，特点是图占比大，每页展示1张
|
||||
("zlib/zlib_17216416", False), # 特殊扫描版3，有的页面是一整张大图，有的页面是通过一条条小图拼起来的，检测图片占比之前需要先按规则把小图拼成大图
|
||||
("the_eye/the_eye_wtl_00023799", False), # 特殊扫描版4，每一页都是一张张小图拼出来的，检测图片占比之前需要先按规则把小图拼成大图
|
||||
("the_eye/the_eye_cdn_00328381", False), # 特殊扫描版5，每一页都是一张张小图拼出来的，存在多个小图多次重复使用情况，检测图片占比之前需要先按规则把小图拼成大图
|
||||
("scihub/scihub_25800000/libgen.scimag25889000-25889999.zip_10.2307/4153991", False), # 特殊扫描版6，只有三页，其中两页是扫描版
|
||||
("scanned_detection/llm-raw-scihub-o.O-0584-8539%2891%2980165-f", False), # 特殊扫描版7，只有一页且由小图拼成大图
|
||||
("scanned_detection/llm-raw-scihub-o.O-bf01427123", False), # 特殊扫描版8，只有3页且全是大图扫描版
|
||||
("scihub/scihub_41200000/libgen.scimag41253000-41253999.zip_10.1080/00222938709460256", False), # 特殊扫描版12，头两页文字版且有一页没图片，后面扫描版11页
|
||||
("scihub/scihub_37000000/libgen.scimag37068000-37068999.zip_10.1080/0015587X.1936.9718622", False) # 特殊扫描版13，头两页文字版且有一页没图片，后面扫描版3页
|
||||
])
|
||||
def test_classify_by_area(book_name, expected_bool_classify_by_area):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
median_width, median_height = get_pdf_page_size_pts(docs)
|
||||
page_width = int(median_width)
|
||||
page_height = int(median_height)
|
||||
img_sz_list = test_data[book_name]["expected_image_info"]
|
||||
total_page = len(docs)
|
||||
text_len_list = get_pdf_textlen_per_page(docs)
|
||||
bool_classify_by_area = classify_by_area(total_page, page_width, page_height, img_sz_list, text_len_list)
|
||||
# assert bool_classify_by_area == expected_bool_classify_by_area
|
||||
|
||||
|
||||
'''
|
||||
广义上的文字版检测，任何一页大于100字，都认为为文字版
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_text_len",
|
||||
[
|
||||
("scihub/scihub_67200000/libgen.scimag67237000-67237999.zip_10.1515/crpm-2017-0020", True), # 文字版，少于50页
|
||||
("scihub/scihub_83300000/libgen.scimag83306000-83306999.zip_10.1007/978-3-658-30153-8", True), # 文字版，多于50页
|
||||
("zhongwenzaixian/zhongwenzaixian_65771414", False), # 完全无字的宣传册
|
||||
])
|
||||
def test_classify_by_text_len(book_name, expected_bool_classify_by_text_len):
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
text_len_list = get_pdf_textlen_per_page(docs)
|
||||
total_page = len(docs)
|
||||
bool_classify_by_text_len = classify_by_text_len(text_len_list, total_page)
|
||||
# assert bool_classify_by_text_len == expected_bool_classify_by_text_len
|
||||
|
||||
|
||||
'''
|
||||
狭义上的文字版检测，需要平均每页字数大于200字
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_avg_words",
|
||||
[
|
||||
("zlib/zlib_21207669", False), # 扫描版，书末尾几页有大纲文字
|
||||
("zlib/zlib_19012845", False), # 扫描版，好几本扫描书的集合，每本书末尾有一页文字页
|
||||
("scihub/scihub_67200000/libgen.scimag67237000-67237999.zip_10.1515/crpm-2017-0020", True),# 正常文字版
|
||||
("zhongwenzaixian/zhongwenzaixian_65771414", False), # 宣传册
|
||||
("zhongwenzaixian/zhongwenzaixian_351879", False), # 图解书/无字or少字
|
||||
("zhongwenzaixian/zhongwenzaixian_61357496_pdfvector", False), # 书法集
|
||||
("zhongwenzaixian/zhongwenzaixian_63684541", False), # 设计图
|
||||
("zhongwenzaixian/zhongwenzaixian_61525978", False), # 绘本
|
||||
("zhongwenzaixian/zhongwenzaixian_63679729", False), # 摄影集
|
||||
|
||||
])
|
||||
def test_classify_by_avg_words(book_name, expected_bool_classify_by_avg_words):
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
text_len_list = get_pdf_textlen_per_page(docs)
|
||||
bool_classify_by_avg_words = classify_by_avg_words(text_len_list)
|
||||
# assert bool_classify_by_avg_words == expected_bool_classify_by_avg_words
|
||||
|
||||
|
||||
'''
|
||||
这个规则只针对特殊扫描版1，因为扫描版1的图片信息都由于junk_list的原因被舍弃了，只能通过图片数量来判断
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_img_num",
|
||||
[
|
||||
("zlib/zlib_21370453", False), # 特殊扫描版1，每页都有所有扫描页图片，特点是图占比大，每页展示1至n张
|
||||
("zlib/zlib_22115997", False), # 特殊扫描版2，类似特1，但是每页数量不完全相等
|
||||
("zlib/zlib_21814957", False), # 特殊扫描版3，类似特1，但是每页数量不完全相等
|
||||
("zlib/zlib_21814955", False), # 特殊扫描版4，类似特1，但是每页数量不完全相等
|
||||
])
|
||||
def test_classify_by_img_num(book_name, expected_bool_classify_by_img_num):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
img_num_list = get_imgs_per_page(docs)
|
||||
img_sz_list = test_data[book_name]["expected_image_info"]
|
||||
bool_classify_by_img_num = classify_by_img_num(img_sz_list, img_num_list)
|
||||
# assert bool_classify_by_img_num == expected_bool_classify_by_img_num
|
||||
|
||||
|
||||
'''
|
||||
排除纵向排版的pdf
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_text_layout",
|
||||
[
|
||||
("vertical_detection/三国演义_繁体竖排版", False), # 竖排版本1
|
||||
("vertical_detection/净空法师_大乘无量寿", False), # 竖排版本2
|
||||
("vertical_detection/om3006239", True), # 横排版本1
|
||||
("vertical_detection/isit.2006.261791", True), # 横排版本2
|
||||
])
|
||||
def test_classify_by_text_layout(book_name, expected_bool_classify_by_text_layout):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
text_layout_per_page = test_data[book_name]["expected_text_layout"]
|
||||
bool_classify_by_text_layout = classify_by_text_layout(text_layout_per_page)
|
||||
# assert bool_classify_by_text_layout == expected_bool_classify_by_text_layout
|
||||
|
||||
|
||||
'''
|
||||
通过检测页面是否由多个窄长条图像组成，来过滤特殊的扫描版
|
||||
这个规则只对窄长条组成的pdf进行识别，而不会识别常规的大图扫描pdf
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_bool_classify_by_img_narrow_strips",
|
||||
[
|
||||
("scihub/scihub_25900000/libgen.scimag25991000-25991999.zip_10.2307/40066695", False), # 特殊扫描版
|
||||
("the_eye/the_eye_wtl_00023799", False), # 特殊扫描版4，每一页都是一张张小图拼出来的，检测图片占比之前需要先按规则把小图拼成大图
|
||||
("the_eye/the_eye_cdn_00328381", False), # 特殊扫描版5，每一页都是一张张小图拼出来的，存在多个小图多次重复使用情况，检测图片占比之前需要先按规则把小图拼成大图
|
||||
("scanned_detection/llm-raw-scihub-o.O-0584-8539%2891%2980165-f", False), # 特殊扫描版7，只有一页且由小图拼成大图
|
||||
("scihub/scihub_25800000/libgen.scimag25889000-25889999.zip_10.2307/4153991", True), # 特殊扫描版6，只有三页，其中两页是扫描版
|
||||
("scanned_detection/llm-raw-scihub-o.O-bf01427123", True), # 特殊扫描版8，只有3页且全是大图扫描版
|
||||
("scihub/scihub_53700000/libgen.scimag53724000-53724999.zip_10.1097/00129191-200509000-00018", True), # 特殊文本版，有一长条，但是只有一条
|
||||
])
|
||||
def test_classify_by_img_narrow_strips(book_name, expected_bool_classify_by_img_narrow_strips):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
img_sz_list = test_data[book_name]["expected_image_info"]
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
median_width, median_height = get_pdf_page_size_pts(docs)
|
||||
page_width = int(median_width)
|
||||
page_height = int(median_height)
|
||||
bool_classify_by_img_narrow_strips = classify_by_img_narrow_strips(page_width, page_height, img_sz_list)
|
||||
# assert bool_classify_by_img_narrow_strips == expected_bool_classify_by_img_narrow_strips
|
||||
@@ -1,80 +0,0 @@
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
|
||||
import fitz
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
|
||||
from magic_pdf.libs.config_reader import get_s3_config_dict
|
||||
|
||||
from magic_pdf.libs.commons import join_path, json_dump_path, read_file, parse_bucket_key
|
||||
from loguru import logger
|
||||
|
||||
test_pdf_dir_path = "s3://llm-pdf-text/unittest/pdf/"
|
||||
|
||||
|
||||
def get_test_pdf_json(book_name):
|
||||
json_path = join_path(json_dump_path, book_name + ".json")
|
||||
s3_config = get_s3_config_dict(json_path)
|
||||
file_content = read_file(json_path, s3_config)
|
||||
json_str = file_content.decode('utf-8')
|
||||
json_object = json.loads(json_str)
|
||||
return json_object
|
||||
|
||||
|
||||
def read_test_file(book_name):
|
||||
test_pdf_path = join_path(test_pdf_dir_path, book_name + ".pdf")
|
||||
s3_config = get_s3_config_dict(test_pdf_path)
|
||||
try:
|
||||
file_content = read_file(test_pdf_path, s3_config)
|
||||
return file_content
|
||||
except Exception as e:
|
||||
if "NoSuchKey" in str(e):
|
||||
logger.warning("File not found in test_pdf_path. Downloading from orig_s3_pdf_path.")
|
||||
try:
|
||||
json_object = get_test_pdf_json(book_name)
|
||||
orig_s3_pdf_path = json_object.get('file_location')
|
||||
s3_config = get_s3_config_dict(orig_s3_pdf_path)
|
||||
file_content = read_file(orig_s3_pdf_path, s3_config)
|
||||
s3_client = get_s3_client(test_pdf_path)
|
||||
bucket_name, bucket_key = parse_bucket_key(test_pdf_path)
|
||||
file_obj = io.BytesIO(file_content)
|
||||
s3_client.upload_fileobj(file_obj, bucket_name, bucket_key)
|
||||
return file_content
|
||||
except Exception as e:
|
||||
logger.exception(e)
|
||||
else:
|
||||
logger.exception(e)
|
||||
|
||||
|
||||
def get_docs_from_test_pdf(book_name):
|
||||
file_content = read_test_file(book_name)
|
||||
return fitz.open("pdf", file_content)
|
||||
|
||||
|
||||
def get_test_json_data(directory_path, json_file_name):
|
||||
with open(os.path.join(directory_path, json_file_name), "r", encoding='utf-8') as f:
|
||||
test_data = json.load(f)
|
||||
return test_data
|
||||
|
||||
|
||||
def get_s3_client(path):
|
||||
s3_config = get_s3_config_dict(path)
|
||||
try:
|
||||
return boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=s3_config["ak"],
|
||||
aws_secret_access_key=s3_config["sk"],
|
||||
endpoint_url=s3_config["endpoint"],
|
||||
config=Config(s3={"addressing_style": "path"}, retries={"max_attempts": 8, "mode": "standard"}),
|
||||
)
|
||||
except:
|
||||
# older boto3 do not support retries.mode param.
|
||||
return boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=s3_config["ak"],
|
||||
aws_secret_access_key=s3_config["sk"],
|
||||
endpoint_url=s3_config["endpoint"],
|
||||
config=Config(s3={"addressing_style": "path"}, retries={"max_attempts": 8}),
|
||||
)
|
||||
@@ -1,84 +0,0 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from magic_pdf.filter.pdf_meta_scan import get_pdf_page_size_pts, get_image_info, get_pdf_text_layout_per_page, get_language
|
||||
from test_commons import get_docs_from_test_pdf, get_test_json_data
|
||||
|
||||
# 获取当前目录
|
||||
current_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
'''
|
||||
获取pdf的宽与高，宽和高各用一个list，分别取中位数
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_width, expected_height",
|
||||
[
|
||||
("zlib/zlib_17058115", 795, 1002), # pdf中最大页与最小页差异极大个例
|
||||
("the_eye/the_eye_wtl_00023799", 616, 785) # 采样的前50页存在中位数大小页面横竖旋转情况
|
||||
])
|
||||
def test_get_pdf_page_size_pts(book_name, expected_width, expected_height):
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
median_width, median_height = get_pdf_page_size_pts(docs)
|
||||
|
||||
# assert int(median_width) == expected_width
|
||||
# assert int(median_height) == expected_height
|
||||
|
||||
|
||||
'''
|
||||
获取pdf前50页的图片信息，为了提速，对特殊扫描版1的情况做了过滤，其余情况都正常取图片信息
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name",
|
||||
[
|
||||
"zlib/zlib_21370453", # 特殊扫描版1，每页都有所有扫描页图片，特点是图占比大，每页展示1至n张
|
||||
"the_eye/the_eye_cdn_00391653", # 特殊文字版1.每页存储所有图片，特点是图片占页面比例不大，每页展示可能为0也可能不止1张，这种pdf需要拿前10页抽样检测img大小和个数，如果符合需要清空junklist
|
||||
"scihub/scihub_08400000/libgen.scimag08489000-08489999.zip_10.1016/0370-1573(90)90070-i", # 扫描版2，每页存储的扫描页图片数量递增，特点是图占比大，每页展示1张，需要清空junklist跑前50页图片信息用于分类判断
|
||||
"zlib/zlib_17216416", # 特殊扫描版3，有的页面是一整张大图，有的页面是通过一条条小图拼起来的
|
||||
"the_eye/the_eye_wtl_00023799", # 特殊扫描版4，每一页都是一张张小图拼出来的
|
||||
"the_eye/the_eye_cdn_00328381", # 特殊扫描版5，每一页都是一张张小图拼出来的，但是存在多个小图多次重复使用情况
|
||||
"scihub/scihub_25800000/libgen.scimag25889000-25889999.zip_10.2307/4153991", # 特殊扫描版6，只有3页且其中两页是扫描页
|
||||
"scanned_detection/llm-raw-scihub-o.O-0584-8539%2891%2980165-f", # 特殊扫描版7，只有一页，且是一张张小图拼出来的
|
||||
"scanned_detection/llm-raw-scihub-o.O-bf01427123", # 特殊扫描版8，只有3页且全是大图扫描版
|
||||
"zlib/zlib_22115997", # 特殊扫描版9，类似特1，但是每页数量不完全相等
|
||||
"zlib/zlib_21814957", # 特殊扫描版10，类似特1，但是每页数量不完全相等
|
||||
"zlib/zlib_21814955", # 特殊扫描版11，类似特1，但是每页数量不完全相等
|
||||
"scihub/scihub_41200000/libgen.scimag41253000-41253999.zip_10.1080/00222938709460256", # 特殊扫描版12，头两页文字版且有一页没图片，后面扫描版11页
|
||||
"scihub/scihub_37000000/libgen.scimag37068000-37068999.zip_10.1080/0015587X.1936.9718622" # 特殊扫描版13，头两页文字版且有一页没图片，后面扫描版3页
|
||||
])
|
||||
def test_get_image_info(book_name):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
page_width_pts, page_height_pts = get_pdf_page_size_pts(docs)
|
||||
image_info, junk_img_bojids = get_image_info(docs, page_width_pts, page_height_pts)
|
||||
|
||||
# assert image_info == test_data[book_name]["expected_image_info"]
|
||||
# assert junk_img_bojids == test_data[book_name]["expected_junk_img_bojids"]
|
||||
|
||||
|
||||
'''
|
||||
获取pdf前50页的文本布局信息，输出list，每个元素为一个页面的横竖排信息
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name",
|
||||
[
|
||||
"vertical_detection/三国演义_繁体竖排版", # 竖排版本1
|
||||
"vertical_detection/净空法师_大乘无量寿", # 竖排版本2
|
||||
"vertical_detection/om3006239", # 横排版本1
|
||||
"vertical_detection/isit.2006.261791" # 横排版本2
|
||||
])
|
||||
def test_get_text_layout_info(book_name):
|
||||
test_data = get_test_json_data(current_directory, "test_metascan_classify_data.json")
|
||||
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
text_layout_info = get_pdf_text_layout_per_page(docs)
|
||||
# assert text_layout_info == test_data[book_name]["expected_text_layout"]
|
||||
|
||||
|
||||
'''
|
||||
获取pdf的语言信息
|
||||
'''
|
||||
@pytest.mark.parametrize("book_name, expected_language",
|
||||
[
|
||||
("scihub/scihub_05000000/libgen.scimag05023000-05023999.zip_10.1034/j.1601-0825.2003.02933.x", "en"), # 英文论文
|
||||
])
|
||||
def test_get_text_language_info(book_name, expected_language):
|
||||
docs = get_docs_from_test_pdf(book_name)
|
||||
text_language = get_language(docs)
|
||||
# assert text_language == expected_language
|
||||
File diff suppressed because one or more lines are too long
@@ -1,687 +0,0 @@
|
||||
[
|
||||
{
|
||||
"layout_dets": [
|
||||
{
|
||||
"category_id": 3,
|
||||
"poly": [
|
||||
776.7277221679688,
|
||||
688.448974609375,
|
||||
1242.224365234375,
|
||||
688.448974609375,
|
||||
1242.224365234375,
|
||||
1182.0628662109375,
|
||||
776.7277221679688,
|
||||
1182.0628662109375
|
||||
],
|
||||
"score": 0.999997079372406
|
||||
},
|
||||
{
|
||||
"category_id": 3,
|
||||
"poly": [
|
||||
775.9269409179688,
|
||||
1389.754638671875,
|
||||
1243.672119140625,
|
||||
1389.754638671875,
|
||||
1243.672119140625,
|
||||
1859.716064453125,
|
||||
775.9269409179688,
|
||||
1859.716064453125
|
||||
],
|
||||
"score": 0.9999949932098389
|
||||
},
|
||||
{
|
||||
"category_id": 1,
|
||||
"poly": [
|
||||
752.11572265625,
|
||||
1939.3634033203125,
|
||||
1430.1146240234375,
|
||||
1939.3634033203125,
|
||||
1430.1146240234375,
|
||||
2041.1771240234375,
|
||||
752.11572265625,
|
||||
2041.1771240234375
|
||||
],
|
||||
"score": 0.999975323677063
|
||||
},
|
||||
{
|
||||
"category_id": 3,
|
||||
"poly": [
|
||||
46.55152893066406,
|
||||
686.12939453125,
|
||||
638.8861083984375,
|
||||
686.12939453125,
|
||||
638.8861083984375,
|
||||
1803.419189453125,
|
||||
46.55152893066406,
|
||||
1803.419189453125
|
||||
],
|
||||
"score": 0.999961256980896
|
||||
},
|
||||
{
|
||||
"category_id": 3,
|
||||
"poly": [
|
||||
33.684722900390625,
|
||||
150.77980041503906,
|
||||
1238.0679931640625,
|
||||
150.77980041503906,
|
||||
1238.0679931640625,
|
||||
524.98291015625,
|
||||
33.684722900390625,
|
||||
524.98291015625
|
||||
],
|
||||
"score": 0.9999504089355469
|
||||
},
|
||||
{
|
||||
"category_id": 1,
|
||||
"poly": [
|
||||
24.685693740844727,
|
||||
1875.9998779296875,
|
||||
703.5064697265625,
|
||||
1875.9998779296875,
|
||||
703.5064697265625,
|
||||
2050.7431640625,
|
||||
24.685693740844727,
|
||||
2050.7431640625
|
||||
],
|
||||
"score": 0.9999105334281921
|
||||
},
|
||||
{
|
||||
"category_id": 1,
|
||||
"poly": [
|
||||
750.97705078125,
|
||||
1252.206787109375,
|
||||
1430.0809326171875,
|
||||
1252.206787109375,
|
||||
1430.0809326171875,
|
||||
1357.2947998046875,
|
||||
750.97705078125,
|
||||
1357.2947998046875
|
||||
],
|
||||
"score": 0.999853789806366
|
||||
},
|
||||
{
|
||||
"category_id": 4,
|
||||
"poly": [
|
||||
904.842041015625,
|
||||
1213.027099609375,
|
||||
1273.5655517578125,
|
||||
1213.027099609375,
|
||||
1273.5655517578125,
|
||||
1242.717529296875,
|
||||
904.842041015625,
|
||||
1242.717529296875
|
||||
],
|
||||
"score": 0.9995817542076111
|
||||
},
|
||||
{
|
||||
"category_id": 4,
|
||||
"poly": [
|
||||
905.3208618164062,
|
||||
1898.5325927734375,
|
||||
1273.1282958984375,
|
||||
1898.5325927734375,
|
||||
1273.1282958984375,
|
||||
1928.9906005859375,
|
||||
905.3208618164062,
|
||||
1928.9906005859375
|
||||
],
|
||||
"score": 0.9986443519592285
|
||||
},
|
||||
{
|
||||
"category_id": 4,
|
||||
"poly": [
|
||||
372.0135498046875,
|
||||
556.02685546875,
|
||||
1084.9647216796875,
|
||||
556.02685546875,
|
||||
1084.9647216796875,
|
||||
586.6792602539062,
|
||||
372.0135498046875,
|
||||
586.6792602539062
|
||||
],
|
||||
"score": 0.9985352754592896
|
||||
},
|
||||
{
|
||||
"category_id": 2,
|
||||
"poly": [
|
||||
1350.63671875,
|
||||
79.77919006347656,
|
||||
1379.6220703125,
|
||||
79.77919006347656,
|
||||
1379.6220703125,
|
||||
99.83788299560547,
|
||||
1350.63671875,
|
||||
99.83788299560547
|
||||
],
|
||||
"score": 0.9973036646842957
|
||||
},
|
||||
{
|
||||
"category_id": 4,
|
||||
"poly": [
|
||||
203.2659912109375,
|
||||
597.2034912109375,
|
||||
1251.0240478515625,
|
||||
597.2034912109375,
|
||||
1251.0240478515625,
|
||||
657.985595703125,
|
||||
203.2659912109375,
|
||||
657.985595703125
|
||||
],
|
||||
"score": 0.9622809886932373
|
||||
},
|
||||
{
|
||||
"category_id": 0,
|
||||
"poly": [
|
||||
70.87332916259766,
|
||||
1834.5714111328125,
|
||||
657.8504638671875,
|
||||
1834.5714111328125,
|
||||
657.8504638671875,
|
||||
1865.07373046875,
|
||||
70.87332916259766,
|
||||
1865.07373046875
|
||||
],
|
||||
"score": 0.8580453395843506
|
||||
},
|
||||
{
|
||||
"category_id": 1,
|
||||
"poly": [
|
||||
189.0360870361328,
|
||||
597.2406616210938,
|
||||
1252.3204345703125,
|
||||
597.2406616210938,
|
||||
1252.3204345703125,
|
||||
658.4781494140625,
|
||||
189.0360870361328,
|
||||
658.4781494140625
|
||||
],
|
||||
"score": 0.3083903193473816
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
1190,
|
||||
1980,
|
||||
1206,
|
||||
1980,
|
||||
1206,
|
||||
1997,
|
||||
1190,
|
||||
1997
|
||||
],
|
||||
"score": 0.51,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
1219,
|
||||
1331,
|
||||
1235,
|
||||
1331,
|
||||
1235,
|
||||
1348,
|
||||
1219,
|
||||
1348
|
||||
],
|
||||
"score": 0.49,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
798,
|
||||
2016,
|
||||
813,
|
||||
2016,
|
||||
813,
|
||||
2033,
|
||||
798,
|
||||
2033
|
||||
],
|
||||
"score": 0.41,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
135,
|
||||
1991,
|
||||
148,
|
||||
1991,
|
||||
148,
|
||||
2006,
|
||||
135,
|
||||
2006
|
||||
],
|
||||
"score": 0.39,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
400,
|
||||
1916,
|
||||
416,
|
||||
1916,
|
||||
416,
|
||||
1933,
|
||||
400,
|
||||
1933
|
||||
],
|
||||
"score": 0.38,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 13,
|
||||
"poly": [
|
||||
1148,
|
||||
1944,
|
||||
1162,
|
||||
1944,
|
||||
1162,
|
||||
1961,
|
||||
1148,
|
||||
1961
|
||||
],
|
||||
"score": 0.31,
|
||||
"latex": ":"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
798.0,
|
||||
1943.0,
|
||||
1147.0,
|
||||
1943.0,
|
||||
1147.0,
|
||||
1968.0,
|
||||
798.0,
|
||||
1968.0
|
||||
],
|
||||
"score": 0.95,
|
||||
"text": "Fig 4 SSCP analysis of FHIT exon 4. T"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
1163.0,
|
||||
1943.0,
|
||||
1425.0,
|
||||
1943.0,
|
||||
1425.0,
|
||||
1968.0,
|
||||
1163.0,
|
||||
1968.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "Tumor tissue ; N :Corresponding"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
755.0,
|
||||
1979.0,
|
||||
1189.0,
|
||||
1979.0,
|
||||
1189.0,
|
||||
2004.0,
|
||||
755.0,
|
||||
2004.0
|
||||
],
|
||||
"score": 0.92,
|
||||
"text": "normal tissue ; M : PBR322/Hae II Marker ; ssDNA"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
1207.0,
|
||||
1979.0,
|
||||
1422.0,
|
||||
1979.0,
|
||||
1422.0,
|
||||
2004.0,
|
||||
1207.0,
|
||||
2004.0
|
||||
],
|
||||
"score": 0.97,
|
||||
"text": "Single-stranded DNA ; ds-"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
755.0,
|
||||
2015.0,
|
||||
797.0,
|
||||
2015.0,
|
||||
797.0,
|
||||
2038.0,
|
||||
755.0,
|
||||
2038.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "DNA"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
814.0,
|
||||
2015.0,
|
||||
996.0,
|
||||
2015.0,
|
||||
996.0,
|
||||
2038.0,
|
||||
814.0,
|
||||
2038.0
|
||||
],
|
||||
"score": 0.98,
|
||||
"text": "Double-stranded DNA"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
71.0,
|
||||
1880.0,
|
||||
698.0,
|
||||
1880.0,
|
||||
698.0,
|
||||
1902.0,
|
||||
71.0,
|
||||
1902.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "Fig 2Alterations of PCR amplified products of FHIT exon 3,4,5 and"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
28.0,
|
||||
1916.0,
|
||||
399.0,
|
||||
1916.0,
|
||||
399.0,
|
||||
1937.0,
|
||||
28.0,
|
||||
1937.0
|
||||
],
|
||||
"score": 0.98,
|
||||
"text": "microsatellite marker D3S1300、D3S1312.A"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
417.0,
|
||||
1916.0,
|
||||
701.0,
|
||||
1916.0,
|
||||
701.0,
|
||||
1937.0,
|
||||
417.0,
|
||||
1937.0
|
||||
],
|
||||
"score": 0.9,
|
||||
"text": "Deletion of exon5(arrows）;B :"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
29.0,
|
||||
1953.0,
|
||||
700.0,
|
||||
1953.0,
|
||||
700.0,
|
||||
1974.0,
|
||||
29.0,
|
||||
1974.0
|
||||
],
|
||||
"score": 0.95,
|
||||
"text": "Deletion of exon 3 A( arrows）;C : Deletion of microsatellite marker D3S1300,"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
28.0,
|
||||
1989.0,
|
||||
134.0,
|
||||
1989.0,
|
||||
134.0,
|
||||
2014.0,
|
||||
28.0,
|
||||
2014.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "D3S1312.T"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
149.0,
|
||||
1989.0,
|
||||
696.0,
|
||||
1989.0,
|
||||
696.0,
|
||||
2014.0,
|
||||
149.0,
|
||||
2014.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "Tumor ; N : Corresponding normal tissue ; L : Corresponding lymph"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
30.0,
|
||||
2027.0,
|
||||
634.0,
|
||||
2027.0,
|
||||
634.0,
|
||||
2047.0,
|
||||
30.0,
|
||||
2047.0
|
||||
],
|
||||
"score": 0.94,
|
||||
"text": "node tissue;M :DL2000 DNA marker;L1:Lewis ;A :A549;S SPAC-1"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
801.0,
|
||||
1259.0,
|
||||
1427.0,
|
||||
1259.0,
|
||||
1427.0,
|
||||
1280.0,
|
||||
801.0,
|
||||
1280.0
|
||||
],
|
||||
"score": 0.94,
|
||||
"text": "Fig 3SSCP analysis of FHIT exon 3.The arrow indicateda deletion of"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
757.0,
|
||||
1294.0,
|
||||
1424.0,
|
||||
1294.0,
|
||||
1424.0,
|
||||
1318.0,
|
||||
757.0,
|
||||
1318.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "exon 3 of 41T. T : Tumor tissue ; N : Corresponding normal tissue ; M PBR322/"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
755.0,
|
||||
1329.0,
|
||||
1218.0,
|
||||
1329.0,
|
||||
1218.0,
|
||||
1355.0,
|
||||
755.0,
|
||||
1355.0
|
||||
],
|
||||
"score": 0.95,
|
||||
"text": "Hae Il Marker / ssDNA : Single-stranded DNA ; dsDNA"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
1236.0,
|
||||
1329.0,
|
||||
1418.0,
|
||||
1329.0,
|
||||
1418.0,
|
||||
1355.0,
|
||||
1236.0,
|
||||
1355.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "Double-strandedDNA"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
910.0,
|
||||
1217.0,
|
||||
1269.0,
|
||||
1217.0,
|
||||
1269.0,
|
||||
1241.0,
|
||||
910.0,
|
||||
1241.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "图3FHIT基因外显子3的SSCP分析"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
909.0,
|
||||
1904.0,
|
||||
1269.0,
|
||||
1904.0,
|
||||
1269.0,
|
||||
1927.0,
|
||||
909.0,
|
||||
1927.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "图4FHIT基因外显子4的SSCP分析"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
374.0,
|
||||
563.0,
|
||||
1077.0,
|
||||
563.0,
|
||||
1077.0,
|
||||
583.0,
|
||||
374.0,
|
||||
583.0
|
||||
],
|
||||
"score": 0.99,
|
||||
"text": "图1FHIT基因外显子3、4、5、8和微卫星灶的PCR扩增产物琼脂糖电泳图"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
1351.0,
|
||||
81.0,
|
||||
1376.0,
|
||||
81.0,
|
||||
1376.0,
|
||||
102.0,
|
||||
1351.0,
|
||||
102.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "13"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
207.0,
|
||||
600.0,
|
||||
1245.0,
|
||||
600.0,
|
||||
1245.0,
|
||||
624.0,
|
||||
207.0,
|
||||
624.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "Fig 1 Agarose electrophoresis of PCR products of exor( A)3 ,4 ,5 ,8 and three microsatellite markers( B）of FHIT gene"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
309.0,
|
||||
634.0,
|
||||
1142.0,
|
||||
634.0,
|
||||
1142.0,
|
||||
662.0,
|
||||
309.0,
|
||||
662.0
|
||||
],
|
||||
"score": 0.97,
|
||||
"text": "M1 :DL2000 DNA marker ; M2 PBR322/Hae Il marker ; T :Tumor ; N :Corresponding normal tissue"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
73.0,
|
||||
1840.0,
|
||||
651.0,
|
||||
1840.0,
|
||||
651.0,
|
||||
1864.0,
|
||||
73.0,
|
||||
1864.0
|
||||
],
|
||||
"score": 1.0,
|
||||
"text": "图2FHIT基因外显子和微卫星灶PCR扩增产物缺失电泳图"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
207.0,
|
||||
600.0,
|
||||
1245.0,
|
||||
600.0,
|
||||
1245.0,
|
||||
625.0,
|
||||
207.0,
|
||||
625.0
|
||||
],
|
||||
"score": 0.96,
|
||||
"text": "Fig 1 Agarose electrophoresis of PCR products of exor A)3 ,4 ,5 ,8 and three microsatellite markers( B)of FHIT gene"
|
||||
},
|
||||
{
|
||||
"category_id": 15,
|
||||
"poly": [
|
||||
309.0,
|
||||
635.0,
|
||||
1142.0,
|
||||
635.0,
|
||||
1142.0,
|
||||
661.0,
|
||||
309.0,
|
||||
661.0
|
||||
],
|
||||
"score": 0.97,
|
||||
"text": "M1 :DL2000 DNA marker ; M2 PBR322/Hae Il marker ; T Tumor ; N :Corresponding normal tissue"
|
||||
}
|
||||
],
|
||||
"page_info": {
|
||||
"page_no": 0,
|
||||
"height": 2080,
|
||||
"width": 1472
|
||||
}
|
||||
}
|
||||
]
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
Binary file not shown.
@@ -1,31 +0,0 @@
|
||||
import json
|
||||
|
||||
from magic_pdf.data.read_api import read_local_pdfs
|
||||
from magic_pdf.model.magic_model import MagicModel
|
||||
|
||||
|
||||
def test_magic_model_image_v2():
|
||||
datasets = read_local_pdfs('tests/unittest/test_model/assets/test_01.pdf')
|
||||
with open('tests/unittest/test_model/assets/test_01.model.json') as f:
|
||||
model_json = json.load(f)
|
||||
|
||||
magic_model = MagicModel(model_json, datasets[0])
|
||||
|
||||
imgs = magic_model.get_imgs_v2(0)
|
||||
print(imgs)
|
||||
|
||||
tables = magic_model.get_tables_v2(0)
|
||||
print(tables)
|
||||
|
||||
|
||||
def test_magic_model_table_v2():
|
||||
datasets = read_local_pdfs('tests/unittest/test_model/assets/test_02.pdf')
|
||||
with open('tests/unittest/test_model/assets/test_02.model.json') as f:
|
||||
model_json = json.load(f)
|
||||
|
||||
magic_model = MagicModel(model_json, datasets[0])
|
||||
tables = magic_model.get_tables_v2(5)
|
||||
print(tables)
|
||||
|
||||
tables = magic_model.get_tables_v2(8)
|
||||
print(tables)
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 58 KiB |
@@ -1,61 +0,0 @@
|
||||
import unittest
|
||||
import os
|
||||
from PIL import Image
|
||||
from lxml import etree
|
||||
|
||||
from magic_pdf.model.sub_modules.model_init import AtomModelSingleton
|
||||
from magic_pdf.model.sub_modules.table.rapidtable.rapid_table import RapidTableModel
|
||||
|
||||
|
||||
class TestppTableModel(unittest.TestCase):
|
||||
def test_image2html(self):
|
||||
img = Image.open(os.path.join(os.path.dirname(__file__), "assets/table.jpg"))
|
||||
atom_model_manager = AtomModelSingleton()
|
||||
ocr_engine = atom_model_manager.get_atom_model(
|
||||
atom_model_name='ocr',
|
||||
ocr_show_log=False,
|
||||
det_db_box_thresh=0.5,
|
||||
det_db_unclip_ratio=1.6,
|
||||
lang='ch'
|
||||
)
|
||||
table_model = RapidTableModel(ocr_engine, 'slanet_plus')
|
||||
html_code, table_cell_bboxes, logic_points, elapse = table_model.predict(img)
|
||||
# 验证生成的 HTML 是否符合预期
|
||||
parser = etree.HTMLParser()
|
||||
tree = etree.fromstring(html_code, parser)
|
||||
|
||||
# 检查 HTML 结构
|
||||
assert tree.find('.//table') is not None, "HTML should contain a <table> element"
|
||||
assert tree.find('.//tr') is not None, "HTML should contain a <tr> element"
|
||||
assert tree.find('.//td') is not None, "HTML should contain a <td> element"
|
||||
|
||||
# 检查具体的表格内容
|
||||
headers = tree.xpath('//table/tr[1]/td')
|
||||
assert len(headers) == 5, "Thead should have 5 columns"
|
||||
assert headers[0].text and headers[0].text.strip() == "Methods", "First header should be 'Methods'"
|
||||
assert headers[1].text and headers[1].text.strip() == "R", "Second header should be 'R'"
|
||||
assert headers[2].text and headers[2].text.strip() == "P", "Third header should be 'P'"
|
||||
assert headers[3].text and headers[3].text.strip() == "F", "Fourth header should be 'F'"
|
||||
assert headers[4].text and headers[4].text.strip() == "FPS", "Fifth header should be 'FPS'"
|
||||
|
||||
# 检查第一行数据
|
||||
first_row = tree.xpath('//table/tr[2]/td')
|
||||
assert len(first_row) == 5, "First row should have 5 cells"
|
||||
assert first_row[0].text and 'SegLink' in first_row[0].text.strip(), "First cell should be 'SegLink [26]'"
|
||||
assert first_row[1].text and first_row[1].text.strip() == "70.0", "Second cell should be '70.0'"
|
||||
assert first_row[2].text and first_row[2].text.strip() == "86.0", "Third cell should be '86.0'"
|
||||
assert first_row[3].text and first_row[3].text.strip() == "77.0", "Fourth cell should be '77.0'"
|
||||
assert first_row[4].text and first_row[4].text.strip() == "8.9", "Fifth cell should be '8.9'"
|
||||
|
||||
# 检查倒数第二行数据
|
||||
second_last_row = tree.xpath('//table/tr[position()=last()-1]/td')
|
||||
assert len(second_last_row) == 5, "second_last_row should have 5 cells"
|
||||
assert second_last_row[0].text and second_last_row[0].text.strip() == "Ours (SynText)", "First cell should be 'Ours (SynText)'"
|
||||
assert second_last_row[1].text and second_last_row[1].text.strip() == "80.68", "Second cell should be '80.68'"
|
||||
assert second_last_row[2].text and second_last_row[2].text.strip() == "85.40", "Third cell should be '85.40'"
|
||||
# assert second_last_row[3].text and second_last_row[3].text.strip() == "82.97", "Fourth cell should be '82.97'"
|
||||
# assert second_last_row[3].text and second_last_row[4].text.strip() == "12.68", "Fifth cell should be '12.68'"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Binary file not shown.
Binary file not shown.
@@ -1,126 +0,0 @@
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from magic_pdf.tools.cli import cli
|
||||
|
||||
|
||||
def test_cli_pdf():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/tools'
|
||||
filename = 'cli_test_01'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir='/tmp/magic_pdf/unittest/tools')
|
||||
|
||||
# run
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
'-p',
|
||||
'tests/unittest/test_tools/assets/cli/pdf/cli_test_01.pdf',
|
||||
'-o',
|
||||
temp_output_dir,
|
||||
],
|
||||
)
|
||||
|
||||
# check
|
||||
assert result.exit_code == 0
|
||||
|
||||
base_output_dir = os.path.join(temp_output_dir, 'cli_test_01/auto')
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 7000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
assert os.path.exists(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.isdir(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.exists(os.path.join(base_output_dir, f'{filename}_content_list.json')) is True
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
|
||||
|
||||
def test_cli_path():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/tools'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir='/tmp/magic_pdf/unittest/tools')
|
||||
|
||||
# run
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli, ['-p', 'tests/unittest/test_tools/assets/cli/path', '-o', temp_output_dir]
|
||||
)
|
||||
|
||||
# check
|
||||
assert result.exit_code == 0
|
||||
|
||||
filename = 'cli_test_01'
|
||||
base_output_dir = os.path.join(temp_output_dir, 'cli_test_01/auto')
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 7000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
assert os.path.exists(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.isdir(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.exists(os.path.join(base_output_dir, f'{filename}_content_list.json')) is True
|
||||
|
||||
base_output_dir = os.path.join(temp_output_dir, 'cli_test_02/auto')
|
||||
filename = 'cli_test_02'
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 5000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
assert os.path.exists(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.isdir(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.exists(os.path.join(base_output_dir, f'{filename}_content_list.json')) is True
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
@@ -1,120 +0,0 @@
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from magic_pdf.tools import cli_dev
|
||||
|
||||
|
||||
def test_cli_pdf():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/tools'
|
||||
filename = 'cli_test_01'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir='/tmp/magic_pdf/unittest/tools')
|
||||
|
||||
# run
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli_dev.cli,
|
||||
[
|
||||
'pdf',
|
||||
'-p',
|
||||
'tests/unittest/test_tools/assets/cli/pdf/cli_test_01.pdf',
|
||||
'-j',
|
||||
'tests/unittest/test_tools/assets/cli_dev/cli_test_01.model.json',
|
||||
'-o',
|
||||
temp_output_dir,
|
||||
],
|
||||
)
|
||||
|
||||
# check
|
||||
assert result.exit_code == 0
|
||||
|
||||
base_output_dir = os.path.join(temp_output_dir, 'cli_test_01/auto')
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_content_list.json'))
|
||||
assert r.st_size > 5000
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 7000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
assert os.path.exists(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.isdir(os.path.join(base_output_dir, 'images')) is True
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
|
||||
|
||||
def test_cli_jsonl():
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/tools'
|
||||
filename = 'cli_test_01'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
temp_output_dir = tempfile.mkdtemp(dir='/tmp/magic_pdf/unittest/tools')
|
||||
|
||||
def mock_read_s3_path(s3path):
|
||||
with open(s3path, 'rb') as f:
|
||||
return f.read()
|
||||
|
||||
cli_dev.read_s3_path = mock_read_s3_path # mock
|
||||
|
||||
# run
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli_dev.cli,
|
||||
[
|
||||
'jsonl',
|
||||
'-j',
|
||||
'tests/unittest/test_tools/assets/cli_dev/cli_test_01.jsonl',
|
||||
'-o',
|
||||
temp_output_dir,
|
||||
],
|
||||
)
|
||||
|
||||
# check
|
||||
assert result.exit_code == 0
|
||||
|
||||
base_output_dir = os.path.join(temp_output_dir, 'cli_test_01/auto')
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_content_list.json'))
|
||||
assert r.st_size > 5000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 7000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
assert os.path.exists(os.path.join(base_output_dir, 'images')) is True
|
||||
assert os.path.isdir(os.path.join(base_output_dir, 'images')) is True
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
@@ -1,59 +0,0 @@
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.tools.common import do_parse
|
||||
|
||||
|
||||
@pytest.mark.parametrize('method', ['auto', 'txt', 'ocr'])
|
||||
def test_common_do_parse(method):
|
||||
import magic_pdf.model as model_config
|
||||
model_config.__use_inside_model__ = True
|
||||
# setup
|
||||
unitest_dir = '/tmp/magic_pdf/unittest/tools'
|
||||
filename = 'fake'
|
||||
os.makedirs(unitest_dir, exist_ok=True)
|
||||
|
||||
temp_output_dir = tempfile.mkdtemp(dir='/tmp/magic_pdf/unittest/tools')
|
||||
|
||||
# run
|
||||
with open('tests/unittest/test_tools/assets/common/cli_test_01.pdf', 'rb') as f:
|
||||
bits = f.read()
|
||||
do_parse(temp_output_dir,
|
||||
filename,
|
||||
bits, [],
|
||||
method,
|
||||
False,
|
||||
f_dump_content_list=True)
|
||||
|
||||
# check
|
||||
base_output_dir = os.path.join(temp_output_dir, f'fake/{method}')
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_content_list.json'))
|
||||
assert r.st_size > 5000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}.md'))
|
||||
assert r.st_size > 7000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_middle.json'))
|
||||
assert r.st_size > 200000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_model.json'))
|
||||
assert r.st_size > 15000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_origin.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_layout.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
r = os.stat(os.path.join(base_output_dir, f'{filename}_spans.pdf'))
|
||||
assert r.st_size > 400000
|
||||
|
||||
os.path.exists(os.path.join(base_output_dir, 'images'))
|
||||
os.path.isdir(os.path.join(base_output_dir, 'images'))
|
||||
|
||||
# teardown
|
||||
shutil.rmtree(temp_output_dir)
|
||||
@@ -1,542 +0,0 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from magic_pdf.libs.boxbase import (__is_overlaps_y_exceeds_threshold,
|
||||
_is_bottom_full_overlap, _is_in,
|
||||
_is_in_or_part_overlap,
|
||||
_is_in_or_part_overlap_with_area_ratio,
|
||||
_is_left_overlap, _is_part_overlap,
|
||||
_is_vertical_full_overlap, _left_intersect,
|
||||
_right_intersect, bbox_distance,
|
||||
bbox_relative_pos, calculate_iou,
|
||||
calculate_overlap_area_2_minbox_area_ratio,
|
||||
calculate_overlap_area_in_bbox1_area_ratio,
|
||||
find_bottom_nearest_text_bbox,
|
||||
find_left_nearest_text_bbox,
|
||||
find_right_nearest_text_bbox,
|
||||
find_top_nearest_text_bbox,
|
||||
get_bbox_in_boundary,
|
||||
get_minbox_if_overlap_by_ratio)
|
||||
from magic_pdf.libs.commons import get_top_percent_list, join_path, mymax
|
||||
from magic_pdf.libs.config_reader import get_s3_config
|
||||
from magic_pdf.libs.path_utils import parse_s3path
|
||||
|
||||
|
||||
# 输入一个列表，如果列表空返回0，否则返回最大元素
|
||||
@pytest.mark.parametrize('list_input, target_num',
|
||||
[
|
||||
([0, 0, 0, 0], 0),
|
||||
([0], 0),
|
||||
([1, 2, 5, 8, 4], 8),
|
||||
([], 0),
|
||||
([1.1, 7.6, 1.009, 9.9], 9.9),
|
||||
([1.0 * 10 ** 2, 3.5 * 10 ** 3, 0.9 * 10 ** 6], 0.9 * 10 ** 6),
|
||||
])
|
||||
def test_list_max(list_input: list, target_num) -> None:
|
||||
"""
|
||||
list_input: 输入列表元素，元素均为数字类型
|
||||
"""
|
||||
assert target_num == mymax(list_input)
|
||||
|
||||
|
||||
# 连接多个参数生成路径信息，使用"/"作为连接符，生成的结果需要是一个合法路径
|
||||
@pytest.mark.parametrize('path_input, target_path', [
|
||||
(['https:', '', 'www.baidu.com'], 'https://www.baidu.com'),
|
||||
(['https:', 'www.baidu.com'], 'https:/www.baidu.com'),
|
||||
(['D:', 'file', 'pythonProject', 'demo' + '.py'], 'D:/file/pythonProject/demo.py'),
|
||||
])
|
||||
def test_join_path(path_input: list, target_path: str) -> None:
|
||||
"""
|
||||
path_input: 输入path的列表，列表元素均为字符串
|
||||
"""
|
||||
assert target_path == join_path(*path_input)
|
||||
|
||||
|
||||
# 获取列表中前百分之多少的元素
|
||||
@pytest.mark.parametrize('num_list, percent, target_num_list', [
|
||||
([], 0.75, []),
|
||||
([-5, -10, 9, 3, 7, -7, 0, 23, -1, -11], 0.8, [23, 9, 7, 3, 0, -1, -5, -7]),
|
||||
([-5, -10, 9, 3, 7, -7, 0, 23, -1, -11], 0, []),
|
||||
([-5, -10, 9, 3, 7, -7, 0, 23, -1, -11, 28], 0.8, [28, 23, 9, 7, 3, 0, -1, -5])
|
||||
])
|
||||
def test_get_top_percent_list(num_list: list, percent: float, target_num_list: list) -> None:
|
||||
"""
|
||||
num_list: 数字列表，列表元素为数字
|
||||
percent: 占比，float, 向下取证
|
||||
"""
|
||||
assert target_num_list == get_top_percent_list(num_list, percent)
|
||||
|
||||
|
||||
# 输入一个s3路径，返回bucket名字和其余部分(key)
|
||||
@pytest.mark.parametrize('s3_path, target_data', [
|
||||
('s3://bucket/path/to/my/file.txt', 'bucket'),
|
||||
('s3a://bucket1/path/to/my/file2.txt', 'bucket1'),
|
||||
# ("/path/to/my/file1.txt", "path"),
|
||||
# ("bucket/path/to/my/file2.txt", "bucket"),
|
||||
])
|
||||
def test_parse_s3path(s3_path: str, target_data: str):
|
||||
"""
|
||||
s3_path: s3路径
|
||||
如果为无效路径，则返回对应的bucket名字和其余部分
|
||||
如果为异常路径 例如：file2.txt，则报异常
|
||||
"""
|
||||
bucket_name, key = parse_s3path(s3_path)
|
||||
assert target_data == bucket_name
|
||||
|
||||
|
||||
# 2个box是否处于包含或者部分重合关系。
|
||||
# 如果某边界重合算重合。
|
||||
# 部分边界重合，其他在内部也算包含
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
((120, 133, 223, 248), (128, 168, 269, 295), True),
|
||||
((137, 53, 245, 157), (134, 11, 200, 147), True), # 部分重合
|
||||
((137, 56, 211, 116), (140, 66, 202, 199), True), # 部分重合
|
||||
((42, 34, 69, 65), (42, 34, 69, 65), True), # 部分重合
|
||||
((39, 63, 87, 106), (37, 66, 85, 109), True), # 部分重合
|
||||
((13, 37, 55, 66), (7, 46, 49, 75), True), # 部分重合
|
||||
((56, 83, 85, 104), (64, 85, 93, 106), True), # 部分重合
|
||||
((12, 53, 48, 94), (14, 53, 50, 94), True), # 部分重合
|
||||
((43, 54, 93, 131), (55, 82, 77, 106), True), # 包含
|
||||
((63, 2, 134, 71), (72, 43, 104, 78), True), # 包含
|
||||
((25, 57, 109, 127), (26, 73, 49, 95), True), # 包含
|
||||
((24, 47, 111, 115), (34, 81, 58, 106), True), # 包含
|
||||
((34, 8, 105, 83), (76, 20, 116, 45), True), # 包含
|
||||
])
|
||||
def test_is_in_or_part_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
"""
|
||||
box1: 坐标数组
|
||||
box2: 坐标数组
|
||||
"""
|
||||
assert target_bool == _is_in_or_part_overlap(box1, box2)
|
||||
|
||||
|
||||
# 如果box1在box2内部，返回True
|
||||
# 如果是部分重合的，则重合面积占box1的比例大于阈值时候返回True
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
((35, 28, 108, 90), (47, 60, 83, 96), False), # 包含 box1 up box2, box2 多半,box1少半
|
||||
((65, 151, 92, 177), (49, 99, 105, 198), True), # 包含 box1 in box2
|
||||
((80, 62, 112, 84), (74, 40, 144, 111), True), # 包含 box1 in box2
|
||||
((65, 88, 127, 144), (92, 102, 131, 139), False), # 包含 box2 多半，box1约一半
|
||||
((92, 102, 131, 139), (65, 88, 127, 144), True), # 包含 box1 多半
|
||||
((100, 93, 199, 168), (169, 126, 198, 165), False), # 包含 box2 in box1
|
||||
((26, 75, 106, 172), (65, 108, 90, 128), False), # 包含 box2 in box1
|
||||
((28, 90, 77, 126), (35, 84, 84, 120), True), # 相交 box1多半，box2多半
|
||||
((37, 6, 69, 52), (28, 3, 60, 49), True), # 相交 box1多半，box2多半
|
||||
((94, 29, 133, 60), (84, 30, 123, 61), True), # 相交 box1多半，box2多半
|
||||
])
|
||||
def test_is_in_or_part_overlap_with_area_ratio(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
out_bool = _is_in_or_part_overlap_with_area_ratio(box1, box2)
|
||||
assert target_bool == out_bool
|
||||
|
||||
|
||||
# box1在box2内部或者box2在box1内部返回True。如果部分边界重合也算作包含。
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
# ((), (), "Error"), # Error
|
||||
((65, 151, 92, 177), (49, 99, 105, 198), True), # 包含 box1 in box2
|
||||
((80, 62, 112, 84), (74, 40, 144, 111), True), # 包含 box1 in box2
|
||||
((76, 140, 154, 277), (121, 326, 192, 384), False), # 分离
|
||||
((65, 88, 127, 144), (92, 102, 131, 139), False), # 包含 box2 多半，box1约一半
|
||||
((92, 102, 131, 139), (65, 88, 127, 144), False), # 包含 box1 多半
|
||||
((68, 94, 118, 120), (68, 90, 118, 122), True), # 包含，box1 in box2 两边x相切
|
||||
((69, 94, 118, 120), (68, 90, 118, 122), True), # 包含，box1 in box2 一边x相切
|
||||
((69, 114, 118, 122), (68, 90, 118, 122), True), # 包含，box1 in box2 一边y相切
|
||||
# ((100, 93, 199, 168), (169, 126, 198, 165), True), # 包含 box2 in box1 Error
|
||||
# ((26, 75, 106, 172), (65, 108, 90, 128), True), # 包含 box2 in box1 Error
|
||||
# ((38, 94, 122, 120), (68, 94, 118, 120), True), # 包含，box2 in box1 两边y相切 Error
|
||||
# ((68, 34, 118, 158), (68, 94, 118, 120), True), # 包含，box2 in box1 两边x相切 Error
|
||||
# ((68, 34, 118, 158), (68, 94, 84, 120), True), # 包含，box2 in box1 一边x相切 Error
|
||||
# ((27, 94, 118, 158), (68, 94, 84, 120), True), # 包含，box2 in box1 一边y相切 Error
|
||||
])
|
||||
def test_is_in(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _is_in(box1, box2)
|
||||
|
||||
|
||||
# 仅仅是部分包含关系，返回True，如果是完全包含关系则返回False
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
((65, 151, 92, 177), (49, 99, 105, 198), False), # 包含 box1 in box2
|
||||
((80, 62, 112, 84), (74, 40, 144, 111), False), # 包含 box1 in box2
|
||||
# ((76, 140, 154, 277), (121, 326, 192, 384), False), # 分离 Error
|
||||
((76, 140, 154, 277), (121, 277, 192, 384), True), # 外相切
|
||||
((65, 88, 127, 144), (92, 102, 131, 139), True), # 包含 box2 多半，box1约一半
|
||||
((92, 102, 131, 139), (65, 88, 127, 144), True), # 包含 box1 多半
|
||||
((68, 94, 118, 120), (68, 90, 118, 122), False), # 包含，box1 in box2 两边x相切
|
||||
((69, 94, 118, 120), (68, 90, 118, 122), False), # 包含，box1 in box2 一边x相切
|
||||
((69, 114, 118, 122), (68, 90, 118, 122), False), # 包含，box1 in box2 一边y相切
|
||||
# ((26, 75, 106, 172), (65, 108, 90, 128), False), # 包含 box2 in box1 Error
|
||||
# ((38, 94, 122, 120), (68, 94, 118, 120), False), # 包含，box2 in box1 两边y相切 Error
|
||||
# ((68, 34, 118, 158), (68, 94, 84, 120), False), # 包含，box2 in box1 一边x相切 Error
|
||||
|
||||
])
|
||||
def test_is_part_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _is_part_overlap(box1, box2)
|
||||
|
||||
|
||||
# left_box右侧是否和right_box左侧有部分重叠
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
(None, None, False),
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), False), # 分离
|
||||
((121, 149, 184, 289), (172, 130, 230, 268), True), # box1 left bottom box2 相交
|
||||
((172, 130, 230, 268), (121, 149, 184, 289), False), # box2 left bottom box1 相交
|
||||
((109, 68, 182, 146), (215, 188, 277, 253), False), # box1 top left box2 分离
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), True), # box1 left top box2 相交
|
||||
((174, 142, 298, 276), (117, 53, 222, 176), False), # box2 left top box1 相交
|
||||
((65, 88, 127, 144), (92, 102, 131, 139), True), # box1 left box2 y:box2 in box1
|
||||
((92, 102, 131, 139), (65, 88, 127, 144), False), # box2 left box1 y:box1 in box2
|
||||
((182, 130, 230, 268), (121, 149, 174, 289), False), # box2 left box1 分离
|
||||
((1, 10, 26, 45), (3, 4, 20, 39), True), # box1 bottom box2 x:box2 in box1
|
||||
])
|
||||
def test_left_intersect(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _left_intersect(box1, box2)
|
||||
|
||||
|
||||
# left_box左侧是否和right_box右侧部分重叠
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
(None, None, False),
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), False), # 分离
|
||||
((121, 149, 184, 289), (172, 130, 230, 268), False), # box1 left bottom box2 相交
|
||||
((172, 130, 230, 268), (121, 149, 184, 289), True), # box2 left bottom box1 相交
|
||||
((109, 68, 182, 146), (215, 188, 277, 253), False), # box1 top left box2 分离
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), False), # box1 left top box2 相交
|
||||
((174, 142, 298, 276), (117, 53, 222, 176), True), # box2 left top box1 相交
|
||||
((65, 88, 127, 144), (92, 102, 131, 139), False), # box1 left box2 y:box2 in box1
|
||||
# ((92, 102, 131, 139), (65, 88, 127, 144), True), # box2 left box1 y:box1 in box2 Error
|
||||
((182, 130, 230, 268), (121, 149, 174, 289), False), # box2 left box1 分离
|
||||
# ((1, 10, 26, 45), (3, 4, 20, 39), False), # box1 bottom box2 x:box2 in box1 Error
|
||||
])
|
||||
def test_right_intersect(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _right_intersect(box1, box2)
|
||||
|
||||
|
||||
# x方向上：要么box1包含box2, 要么box2包含box1。不能部分包含
|
||||
# y方向上：box1和box2有重叠
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
# (None, None, False), # Error
|
||||
((35, 28, 108, 90), (47, 60, 83, 96), True), # box1 top box2, x:box2 in box1, y:有重叠
|
||||
((35, 28, 98, 90), (27, 60, 103, 96), True), # box1 top box2, x:box1 in box2, y:有重叠
|
||||
((57, 77, 130, 210), (59, 219, 119, 293), False), # box1 top box2, x: box2 in box1, y:无重叠
|
||||
((47, 60, 83, 96), (35, 28, 108, 90), True), # box2 top box1, x:box1 in box2, y:有重叠
|
||||
((27, 60, 103, 96), (35, 28, 98, 90), True), # box2 top box1, x:box2 in box1, y:有重叠
|
||||
((59, 219, 119, 293), (57, 77, 130, 210), False), # box2 top box1, x: box1 in box2, y:无重叠
|
||||
((35, 28, 55, 90), (57, 60, 83, 96), False), # box1 top box2, x:无重叠, y:有重叠
|
||||
((47, 60, 63, 96), (65, 28, 108, 90), False), # box2 top box1, x:无重叠, y:有重叠
|
||||
])
|
||||
def test_is_vertical_full_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _is_vertical_full_overlap(box1, box2)
|
||||
|
||||
|
||||
# 检查box1下方和box2的上方有轻微的重叠，轻微程度收到y_tolerance的限制
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
(None, None, False),
|
||||
((35, 28, 108, 90), (47, 89, 83, 116), True), # box1 top box2, y:有重叠
|
||||
((35, 28, 108, 90), (47, 60, 83, 96), False), # box1 top box2, y:有重叠且过多
|
||||
((57, 77, 130, 210), (59, 219, 119, 293), False), # box1 top box2, y:无重叠
|
||||
((47, 60, 83, 96), (35, 28, 108, 90), False), # box2 top box1, y:有重叠且过多
|
||||
((27, 89, 103, 116), (35, 28, 98, 90), False), # box2 top box1, y:有重叠
|
||||
((59, 219, 119, 293), (57, 77, 130, 210), False), # box2 top box1, y:无重叠
|
||||
])
|
||||
def test_is_bottom_full_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _is_bottom_full_overlap(box1, box2)
|
||||
|
||||
|
||||
# 检查box1的左侧是否和box2有重叠
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
(None, None, False),
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), False), # 分离
|
||||
# ((121, 149, 184, 289), (172, 130, 230, 268), False), # box1 left bottom box2 相交 Error
|
||||
# ((172, 130, 230, 268), (121, 149, 184, 289), True), # box2 left bottom box1 相交 Error
|
||||
((109, 68, 182, 146), (215, 188, 277, 253), False), # box1 top left box2 分离
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), False), # box1 left top box2 相交
|
||||
# ((174, 142, 298, 276), (117, 53, 222, 176), True), # box2 left top box1 相交 Error
|
||||
# ((65, 88, 127, 144), (92, 102, 131, 139), False), # box1 left box2 y:box2 in box1 Error
|
||||
((1, 10, 26, 45), (3, 4, 20, 39), True), # box1 middle bottom box2 x:box2 in box1
|
||||
|
||||
])
|
||||
def test_is_left_overlap(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == _is_left_overlap(box1, box2)
|
||||
|
||||
|
||||
# 查两个bbox在y轴上是否有重叠，并且该重叠区域的高度占两个bbox高度更低的那个超过阈值
|
||||
@pytest.mark.parametrize('box1, box2, target_bool', [
|
||||
# (None, None, "Error"), # Error
|
||||
((51, 69, 192, 147), (75, 48, 132, 187), True), # y: box1 in box2
|
||||
((51, 39, 192, 197), (75, 48, 132, 187), True), # y: box2 in box1
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), False), # y: box1 top box2
|
||||
((109, 68, 182, 196), (215, 188, 277, 253), False), # y: box1 top box2 little
|
||||
((109, 68, 182, 196), (215, 78, 277, 253), True), # y: box1 top box2 more
|
||||
((109, 68, 182, 196), (215, 138, 277, 213), False), # y: box1 top box2 more but lower overlap_ratio_threshold
|
||||
((109, 68, 182, 196), (215, 138, 277, 203), True), # y: box1 top box2 more and more overlap_ratio_threshold
|
||||
])
|
||||
def test_is_overlaps_y_exceeds_threshold(box1: tuple, box2: tuple, target_bool: bool) -> None:
|
||||
assert target_bool == __is_overlaps_y_exceeds_threshold(box1, box2)
|
||||
|
||||
|
||||
# Determine the coordinates of the intersection rectangle
|
||||
@pytest.mark.parametrize('box1, box2, target_num', [
|
||||
# (None, None, "Error"), # Error
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), 0.0), # 分离
|
||||
((76, 140, 154, 277), (121, 326, 192, 384), 0.0), # 分离
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), 0.0), # 分离
|
||||
((109, 68, 182, 196), (175, 138, 277, 213), 0.024475524475524476), # 相交
|
||||
((56, 90, 170, 219), (103, 212, 171, 304), 0.02288586346557361), # 相交
|
||||
((109, 126, 204, 245), (130, 127, 232, 186), 0.33696071621517326), # 相交
|
||||
((109, 126, 204, 245), (110, 127, 232, 206), 0.5493822593770807), # 相交
|
||||
((76, 140, 154, 277), (121, 277, 192, 384), 0.0) # 相切
|
||||
])
|
||||
def test_calculate_iou(box1: tuple, box2: tuple, target_num: float) -> None:
|
||||
assert target_num == calculate_iou(box1, box2)
|
||||
|
||||
|
||||
# 计算box1和box2的重叠面积占最小面积的box的比例
|
||||
@pytest.mark.parametrize('box1, box2, target_num', [
|
||||
# (None, None, "Error"), # Error
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), 0.0), # 分离
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), 0.0), # 分离
|
||||
((76, 140, 154, 277), (121, 326, 192, 384), 0.0), # 分离
|
||||
((76, 140, 154, 277), (121, 277, 192, 384), 0.0), # 相切
|
||||
((109, 126, 204, 245), (110, 127, 232, 206), 0.7704918032786885), # 相交
|
||||
((56, 90, 170, 219), (103, 212, 171, 304), 0.07496803069053709), # 相交
|
||||
((121, 149, 184, 289), (172, 130, 230, 268), 0.17841079460269865), # 相交
|
||||
((51, 69, 192, 147), (75, 48, 132, 187), 0.5611510791366906), # 相交
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), 0.12636469221835075), # 相交
|
||||
((102, 60, 233, 203), (70, 190, 220, 319), 0.08188757807078417), # 相交
|
||||
((109, 126, 204, 245), (130, 127, 232, 186), 0.7254901960784313), # 相交
|
||||
])
|
||||
def test_calculate_overlap_area_2_minbox_area_ratio(box1: tuple, box2: tuple, target_num: float) -> None:
|
||||
assert target_num == calculate_overlap_area_2_minbox_area_ratio(box1, box2)
|
||||
|
||||
|
||||
# 计算box1和box2的重叠面积占bbox1的比例
|
||||
@pytest.mark.parametrize('box1, box2, target_num', [
|
||||
# (None, None, "Error"), # Error
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), 0.0), # 分离
|
||||
((88, 81, 222, 173), (60, 221, 123, 358), 0.0), # 分离
|
||||
((76, 140, 154, 277), (121, 326, 192, 384), 0.0), # 分离
|
||||
((76, 140, 154, 277), (121, 277, 192, 384), 0.0), # 相切
|
||||
((142, 109, 238, 164), (134, 164, 224, 270), 0.0), # 相切
|
||||
((109, 126, 204, 245), (110, 127, 232, 206), 0.6568774878372402), # 相交
|
||||
((56, 90, 170, 219), (103, 212, 171, 304), 0.03189174486604107), # 相交
|
||||
((121, 149, 184, 289), (172, 130, 230, 268), 0.1619047619047619), # 相交
|
||||
((51, 69, 192, 147), (75, 48, 132, 187), 0.40425531914893614), # 相交
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), 0.12636469221835075), # 相交
|
||||
((102, 60, 233, 203), (70, 190, 220, 319), 0.08188757807078417), # 相交
|
||||
((109, 126, 204, 245), (130, 127, 232, 186), 0.38620079610791685), # 相交
|
||||
])
|
||||
def test_calculate_overlap_area_in_bbox1_area_ratio(box1: tuple, box2: tuple, target_num: float) -> None:
|
||||
assert target_num == calculate_overlap_area_in_bbox1_area_ratio(box1, box2)
|
||||
|
||||
|
||||
# 计算两个bbox重叠的面积占最小面积的box的比例，如果比例大于ratio，则返回小的那个bbox,否则返回None
|
||||
@pytest.mark.parametrize('box1, box2, ratio, target_box', [
|
||||
# (None, None, 0.8, "Error"), # Error
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), 0.0, None), # 分离
|
||||
((109, 126, 204, 245), (110, 127, 232, 206), 0.5, (110, 127, 232, 206)),
|
||||
((56, 90, 170, 219), (103, 212, 171, 304), 0.5, None),
|
||||
((121, 149, 184, 289), (172, 130, 230, 268), 0.5, None),
|
||||
((51, 69, 192, 147), (75, 48, 132, 187), 0.5, (75, 48, 132, 187)),
|
||||
((117, 53, 222, 176), (174, 142, 298, 276), 0.5, None),
|
||||
((102, 60, 233, 203), (70, 190, 220, 319), 0.5, None),
|
||||
((109, 126, 204, 245), (130, 127, 232, 186), 0.5, (130, 127, 232, 186)),
|
||||
])
|
||||
def test_get_minbox_if_overlap_by_ratio(box1: tuple, box2: tuple, ratio: float, target_box: list) -> None:
|
||||
assert target_box == get_minbox_if_overlap_by_ratio(box1, box2, ratio)
|
||||
|
||||
|
||||
# 根据boundry获取在这个范围内的所有的box的列表，完全包含关系
|
||||
@pytest.mark.parametrize('boxes, boundary, target_boxs', [
|
||||
# ([], (), "Error"), # Error
|
||||
([], (110, 340, 209, 387), []),
|
||||
([(142, 109, 238, 164)], (134, 211, 224, 270), []), # 分离
|
||||
([(109, 126, 204, 245), (110, 127, 232, 206)], (105, 116, 258, 300), [(109, 126, 204, 245), (110, 127, 232, 206)]),
|
||||
([(109, 126, 204, 245), (110, 127, 232, 206)], (105, 116, 258, 230), [(110, 127, 232, 206)]),
|
||||
([(81, 280, 123, 315), (282, 203, 342, 247), (183, 100, 300, 155), (46, 99, 133, 148), (33, 156, 97, 211),
|
||||
(137, 29, 287, 87)], (80, 90, 249, 200), []),
|
||||
([(81, 280, 123, 315), (282, 203, 342, 247), (183, 100, 300, 155), (46, 99, 133, 148), (33, 156, 97, 211),
|
||||
(137, 29, 287, 87)], (30, 20, 349, 320),
|
||||
[(81, 280, 123, 315), (282, 203, 342, 247), (183, 100, 300, 155), (46, 99, 133, 148), (33, 156, 97, 211),
|
||||
(137, 29, 287, 87)]),
|
||||
([(81, 280, 123, 315), (282, 203, 342, 247), (183, 100, 300, 155), (46, 99, 133, 148), (33, 156, 97, 211),
|
||||
(137, 29, 287, 87)], (30, 20, 200, 320),
|
||||
[(81, 280, 123, 315), (46, 99, 133, 148), (33, 156, 97, 211)]),
|
||||
])
|
||||
def test_get_bbox_in_boundary(boxes: list, boundary: tuple, target_boxs: list) -> None:
|
||||
assert target_boxs == get_bbox_in_boundary(boxes, boundary)
|
||||
|
||||
|
||||
# 寻找上方距离最近的box,margin 4个单位， x方向有重合，y方向最近的
|
||||
@pytest.mark.parametrize('pymu_blocks, obj_box, target_boxs', [
|
||||
([{'bbox': (81, 280, 123, 315)}, {'bbox': (282, 203, 342, 247)}, {'bbox': (183, 100, 300, 155)},
|
||||
{'bbox': (46, 99, 133, 148)}, {'bbox': (33, 156, 97, 211)},
|
||||
{'bbox': (137, 29, 287, 87)}], (81, 280, 123, 315), {'bbox': (33, 156, 97, 211)}),
|
||||
# ([{"bbox": (168, 120, 263, 159)},
|
||||
# {"bbox": (231, 61, 279, 159)},
|
||||
# {"bbox": (35, 85, 136, 110)},
|
||||
# {"bbox": (228, 193, 347, 225)},
|
||||
# {"bbox": (144, 264, 188, 323)},
|
||||
# {"bbox": (62, 37, 126, 64)}], (228, 193, 347, 225),
|
||||
# [{"bbox": (168, 120, 263, 159)}, {"bbox": (231, 61, 279, 159)}]), # y：方向最近的有两个，x: 两个均有重合 Error
|
||||
([{'bbox': (35, 85, 136, 159)},
|
||||
{'bbox': (168, 120, 263, 159)},
|
||||
{'bbox': (231, 61, 279, 118)},
|
||||
{'bbox': (228, 193, 347, 225)},
|
||||
{'bbox': (144, 264, 188, 323)},
|
||||
{'bbox': (62, 37, 126, 64)}], (228, 193, 347, 225),
|
||||
{'bbox': (168, 120, 263, 159)},), # y:方向最近的有两个，x:只有一个有重合
|
||||
([{'bbox': (239, 115, 379, 167)},
|
||||
{'bbox': (33, 237, 104, 262)},
|
||||
{'bbox': (124, 288, 168, 325)},
|
||||
{'bbox': (242, 291, 379, 340)},
|
||||
{'bbox': (55, 117, 121, 154)},
|
||||
{'bbox': (266, 183, 384, 217)}, ], (124, 288, 168, 325), {'bbox': (55, 117, 121, 154)}),
|
||||
([{'bbox': (239, 115, 379, 167)},
|
||||
{'bbox': (33, 237, 104, 262)},
|
||||
{'bbox': (124, 288, 168, 325)},
|
||||
{'bbox': (242, 291, 379, 340)},
|
||||
{'bbox': (55, 117, 119, 154)},
|
||||
{'bbox': (266, 183, 384, 217)}, ], (124, 288, 168, 325), None), # x没有重合
|
||||
([{'bbox': (80, 90, 249, 200)},
|
||||
{'bbox': (183, 100, 240, 155)}, ], (183, 100, 240, 155), None), # 包含
|
||||
])
|
||||
def test_find_top_nearest_text_bbox(pymu_blocks: list, obj_box: tuple, target_boxs: dict) -> None:
|
||||
assert target_boxs == find_top_nearest_text_bbox(pymu_blocks, obj_box)
|
||||
|
||||
|
||||
# 寻找下方距离自己最近的box, x方向有重合，y方向最近的
|
||||
@pytest.mark.parametrize('pymu_blocks, obj_box, target_boxs', [
|
||||
([{'bbox': (165, 96, 300, 114)},
|
||||
{'bbox': (11, 157, 139, 201)},
|
||||
{'bbox': (124, 208, 265, 262)},
|
||||
{'bbox': (124, 283, 248, 306)},
|
||||
{'bbox': (39, 267, 84, 301)},
|
||||
{'bbox': (36, 89, 114, 145)}, ], (165, 96, 300, 114), {'bbox': (124, 208, 265, 262)}),
|
||||
([{'bbox': (187, 37, 303, 49)},
|
||||
{'bbox': (2, 227, 90, 283)},
|
||||
{'bbox': (158, 174, 200, 212)},
|
||||
{'bbox': (259, 174, 324, 228)},
|
||||
{'bbox': (205, 61, 316, 97)},
|
||||
{'bbox': (295, 248, 374, 287)}, ], (205, 61, 316, 97), {'bbox': (259, 174, 324, 228)}), # y有两个最近的, x只有一个重合
|
||||
# ([{"bbox": (187, 37, 303, 49)},
|
||||
# {"bbox": (2, 227, 90, 283)},
|
||||
# {"bbox": (259, 174, 324, 228)},
|
||||
# {"bbox": (205, 61, 316, 97)},
|
||||
# {"bbox": (295, 248, 374, 287)},
|
||||
# {"bbox": (158, 174, 209, 212)}, ], (205, 61, 316, 97),
|
||||
# [{"bbox": (259, 174, 324, 228)}, {"bbox": (158, 174, 209, 212)}]), # x有重合，y有两个最近的 Error
|
||||
([{'bbox': (287, 132, 398, 191)},
|
||||
{'bbox': (44, 141, 163, 188)},
|
||||
{'bbox': (132, 191, 240, 241)},
|
||||
{'bbox': (81, 25, 142, 67)},
|
||||
{'bbox': (74, 297, 116, 314)},
|
||||
{'bbox': (77, 84, 224, 107)}, ], (287, 132, 398, 191), None), # x没有重合
|
||||
([{'bbox': (80, 90, 249, 200)},
|
||||
{'bbox': (183, 100, 240, 155)}, ], (183, 100, 240, 155), None), # 包含
|
||||
])
|
||||
def test_find_bottom_nearest_text_bbox(pymu_blocks: list, obj_box: tuple, target_boxs: dict) -> None:
|
||||
assert target_boxs == find_bottom_nearest_text_bbox(pymu_blocks, obj_box)
|
||||
|
||||
|
||||
# 寻找左侧距离自己最近的box, y方向有重叠，x方向最近
|
||||
@pytest.mark.parametrize('pymu_blocks, obj_box, target_boxs', [
|
||||
([{'bbox': (80, 90, 249, 200)}, {'bbox': (183, 100, 240, 155)}], (183, 100, 240, 155), None), # 包含
|
||||
([{'bbox': (28, 90, 77, 126)}, {'bbox': (35, 84, 84, 120)}], (35, 84, 84, 120), None), # y:重叠，x:重叠大于2
|
||||
([{'bbox': (28, 90, 77, 126)}, {'bbox': (75, 84, 134, 120)}], (75, 84, 134, 120), {'bbox': (28, 90, 77, 126)}),
|
||||
# y:重叠，x:重叠小于等于2
|
||||
([{'bbox': (239, 115, 379, 167)},
|
||||
{'bbox': (33, 237, 104, 262)},
|
||||
{'bbox': (124, 288, 168, 325)},
|
||||
{'bbox': (242, 291, 379, 340)},
|
||||
{'bbox': (55, 113, 161, 154)},
|
||||
{'bbox': (266, 123, 384, 217)}], (266, 123, 384, 217), {'bbox': (55, 113, 161, 154)}), # y重叠，x left
|
||||
# ([{"bbox": (136, 219, 268, 240)},
|
||||
# {"bbox": (169, 115, 268, 181)},
|
||||
# {"bbox": (33, 237, 104, 262)},
|
||||
# {"bbox": (124, 288, 168, 325)},
|
||||
# {"bbox": (55, 117, 161, 154)},
|
||||
# {"bbox": (266, 183, 384, 217)}], (266, 183, 384, 217),
|
||||
# [{"bbox": (136, 219, 267, 240)}, {"bbox": (169, 115, 267, 181)}]), # y有重叠，x重叠小于2或者在left Error
|
||||
])
|
||||
def test_find_left_nearest_text_bbox(pymu_blocks: list, obj_box: tuple, target_boxs: dict) -> None:
|
||||
assert target_boxs == find_left_nearest_text_bbox(pymu_blocks, obj_box)
|
||||
|
||||
|
||||
# 寻找右侧距离自己最近的box, y方向有重叠，x方向最近
|
||||
@pytest.mark.parametrize('pymu_blocks, obj_box, target_boxs', [
|
||||
([{'bbox': (80, 90, 249, 200)}, {'bbox': (183, 100, 240, 155)}], (183, 100, 240, 155), None), # 包含
|
||||
([{'bbox': (28, 90, 77, 126)}, {'bbox': (35, 84, 84, 120)}], (28, 90, 77, 126), None), # y:重叠，x:重叠大于2
|
||||
([{'bbox': (28, 90, 77, 126)}, {'bbox': (75, 84, 134, 120)}], (28, 90, 77, 126), {'bbox': (75, 84, 134, 120)}),
|
||||
# y:重叠，x:重叠小于等于2
|
||||
([{'bbox': (239, 115, 379, 167)},
|
||||
{'bbox': (33, 237, 104, 262)},
|
||||
{'bbox': (124, 288, 168, 325)},
|
||||
{'bbox': (242, 291, 379, 340)},
|
||||
{'bbox': (55, 113, 161, 154)},
|
||||
{'bbox': (266, 123, 384, 217)}], (55, 113, 161, 154), {'bbox': (239, 115, 379, 167)}), # y重叠，x right
|
||||
# ([{"bbox": (169, 115, 298, 181)},
|
||||
# {"bbox": (169, 219, 268, 240)},
|
||||
# {"bbox": (33, 177, 104, 262)},
|
||||
# {"bbox": (124, 288, 168, 325)},
|
||||
# {"bbox": (55, 117, 161, 154)},
|
||||
# {"bbox": (266, 183, 384, 217)}], (33, 177, 104, 262),
|
||||
# [{"bbox": (169, 115, 298, 181)}, {"bbox": (169, 219, 268, 240)}]), # y有重叠，x重叠小于2或者在right Error
|
||||
])
|
||||
def test_find_right_nearest_text_bbox(pymu_blocks: list, obj_box: tuple, target_boxs: dict) -> None:
|
||||
assert target_boxs == find_right_nearest_text_bbox(pymu_blocks, obj_box)
|
||||
|
||||
|
||||
# 判断两个矩形框的相对位置关系 (left, right, bottom, top)
|
||||
@pytest.mark.parametrize('box1, box2, target_box', [
|
||||
# (None, None, "Error"), # Error
|
||||
((80, 90, 249, 200), (183, 100, 240, 155), (False, False, False, False)), # 包含
|
||||
# ((124, 81, 222, 173), (60, 221, 123, 358), (False, True, False, True)), # 分离，右上 Error
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), (False, False, False, True)), # 分离，上
|
||||
# ((51, 69, 192, 147), (205, 198, 282, 297), (True, False, False, True)), # 分离，左上 Error
|
||||
# ((101, 149, 164, 289), (172, 130, 230, 268), (True, False, False, False)), # 分离，左 Error
|
||||
# ((69, 196, 124, 285), (130, 127, 232, 186), (True, False, True, False)), # 分离，左下 Error
|
||||
((103, 212, 171, 304), (56, 90, 170, 209), (False, False, True, False)), # 分离，下
|
||||
# ((124, 367, 222, 415), (60, 221, 123, 358), (False, True, True, False)), # 分离，右下 Error
|
||||
# ((172, 130, 230, 268), (101, 149, 164, 289), (False, True, False, False)), # 分离，右 Error
|
||||
])
|
||||
def test_bbox_relative_pos(box1: tuple, box2: tuple, target_box: tuple) -> None:
|
||||
assert target_box == bbox_relative_pos(box1, box2)
|
||||
|
||||
|
||||
# 计算两个矩形框的距离
|
||||
"""
|
||||
受bbox_relative_pos方法的影响，左右相反，这里计算结果全部受影响，在错误的基础上计算出了正确的结果
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize('box1, box2, target_num', [
|
||||
# (None, None, "Error"), # Error
|
||||
((80, 90, 249, 200), (183, 100, 240, 155), 0.0), # 包含
|
||||
((142, 109, 238, 164), (134, 211, 224, 270), 47.0), # 分离，上
|
||||
((103, 212, 171, 304), (56, 90, 170, 209), 3.0), # 分离，下
|
||||
((101, 149, 164, 289), (172, 130, 230, 268), 8.0), # 分离，左
|
||||
((172, 130, 230, 268), (101, 149, 164, 289), 8.0), # 分离，右
|
||||
((80.3, 90.8, 249.0, 200.5), (183.8, 100.6, 240.2, 155.1), 0.0), # 包含
|
||||
((142.3, 109.5, 238.9, 164.2), (134.4, 211.2, 224.8, 270.1), 47.0), # 分离，上
|
||||
((103.5, 212.6, 171.1, 304.8), (56.1, 90.9, 170.6, 209.2), 3.4), # 分离，下
|
||||
((101.1, 149.3, 164.9, 289.0), (172.1, 130.1, 230.5, 268.5), 7.2), # 分离，左
|
||||
((172.1, 130.3, 230.1, 268.1), (101.2, 149.9, 164.3, 289.1), 7.8), # 分离，右
|
||||
((124.3, 81.1, 222.5, 173.8), (60.3, 221.5, 123.0, 358.9), 47.717711596429254), # 分离，右上
|
||||
((51.2, 69.31, 192.5, 147.9), (205.0, 198.1, 282.98, 297.09), 51.73287156151299), # 分离，左上
|
||||
((124.3, 367.1, 222.9, 415.7), (60.9, 221.4, 123.2, 358.6), 8.570880934886448), # 分离，右下
|
||||
((69.9, 196.2, 124.1, 285.7), (130.0, 127.3, 232.6, 186.1), 11.69700816448377), # 分离，左下
|
||||
])
|
||||
def test_bbox_distance(box1: tuple, box2: tuple, target_num: float) -> None:
|
||||
assert target_num - bbox_distance(box1, box2) < 1
|
||||
|
||||
|
||||
@pytest.mark.skip(reason='skip')
|
||||
# 根据bucket_name获取s3配置ak,sk,endpoint
|
||||
def test_get_s3_config() -> None:
|
||||
bucket_name = os.getenv('bucket_name')
|
||||
target_data = os.getenv('target_data')
|
||||
assert convert_string_to_list(target_data) == list(get_s3_config(bucket_name))
|
||||
|
||||
|
||||
def convert_string_to_list(s):
|
||||
cleaned_s = s.strip("'")
|
||||
items = cleaned_s.split(',')
|
||||
cleaned_items = [item.strip() for item in items]
|
||||
return cleaned_items
|
||||
Reference in New Issue
Block a user