mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-10-25 17:42:08 +08:00
250 lines
11 KiB
Python
250 lines
11 KiB
Python
# Copyright (c) Opendatalab. All rights reserved.
|
|
"""验证 ONNX 检测真批处理：真实概率图/几何对照和带调用记录的 PDF 解析。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import resource
|
|
import sys
|
|
import time
|
|
from importlib.metadata import version
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from mineru.model.runtime.onnx import ort_session
|
|
from .validate_onnx_models import overlay, parse_documents, write_json
|
|
|
|
__all__ = ["RecordingSession", "components", "integration"]
|
|
|
|
|
|
class RecordingSession:
|
|
"""只在验证脚本记录 ORT 的真实输入批次，不更改模型图或生产调度。"""
|
|
|
|
def __init__(self, session: Any, name: str, *, retain_outputs: bool = False) -> None:
|
|
"""保存原会话；组件对照可额外保留概率图。"""
|
|
self.session = session
|
|
self.name = name
|
|
self.retain_outputs = retain_outputs
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.outputs: list[list[np.ndarray]] = []
|
|
|
|
def get_inputs(self) -> list[Any]:
|
|
"""透传图输入签名。"""
|
|
return self.session.get_inputs()
|
|
|
|
def get_outputs(self) -> list[Any]:
|
|
"""透传图输出签名。"""
|
|
return self.session.get_outputs()
|
|
|
|
def run(self, output_names: list[str] | None, input_feed: dict[str, np.ndarray]) -> list[np.ndarray]:
|
|
"""记录送入 ORT 的形状和纯推理耗时。"""
|
|
started = time.perf_counter()
|
|
result = self.session.run(output_names, input_feed)
|
|
self.calls.append(
|
|
{"inputs": {key: list(value.shape) for key, value in input_feed.items()}, "seconds": time.perf_counter() - started}
|
|
)
|
|
if self.retain_outputs:
|
|
self.outputs.append(result)
|
|
return result
|
|
|
|
|
|
def components(args: argparse.Namespace) -> dict[str, Any]:
|
|
"""用不同页面、翻转页、空图和真实印章比较逐项与批量输出。"""
|
|
from mineru.model.layout.pp_doclayout_v2_onnx import PPDocLayoutV2LayoutModelONNX
|
|
from mineru.model.ocr.pp_ocr_v6_onnx import TextDetectorONNX
|
|
|
|
repo = args.models_dir / "MinerU-4_models_onnx"
|
|
pages = [
|
|
Image.open(args.samples_dir / name).convert("RGB") for name in ("demo1-page1.png", "demo2-page3.png", "seal-page.png")
|
|
]
|
|
layout = PPDocLayoutV2LayoutModelONNX(
|
|
str(repo / "Layout/PP-DocLayoutV2/inference.onnx"), config_path=str(repo / "Layout/PP-DocLayoutV2/inference.yml")
|
|
)
|
|
observer = RecordingSession(layout.session, "layout")
|
|
layout.session = observer
|
|
layout.predict(pages[0])
|
|
observer.calls.clear()
|
|
images = [*pages, Image.new("RGB", (800, 1000), "white")]
|
|
started = time.perf_counter()
|
|
single = layout.batch_predict(images, batch_size=1)
|
|
single_seconds = time.perf_counter() - started
|
|
single_calls = list(observer.calls)
|
|
observer.calls.clear()
|
|
started = time.perf_counter()
|
|
batch = layout.batch_predict(images, batch_size=2)
|
|
batch_seconds = time.perf_counter() - started
|
|
checks = []
|
|
for index, (left, right) in enumerate(zip(single, batch)):
|
|
structure_match = [(x["label"], x["index"]) for x in left] == [(x["label"], x["index"]) for x in right]
|
|
checks.append(
|
|
{
|
|
"index": index,
|
|
"boxes_single": len(left),
|
|
"boxes_batch": len(right),
|
|
"labels_order_match": structure_match,
|
|
"exact_match": left == right,
|
|
"max_bbox_delta": max(
|
|
(float(np.abs(np.array(a["bbox"]) - b["bbox"]).max()) for a, b in zip(left, right)), default=0
|
|
)
|
|
if structure_match
|
|
else None,
|
|
"max_score_delta": max((abs(a["score"] - b["score"]) for a, b in zip(left, right)), default=0)
|
|
if structure_match
|
|
else None,
|
|
}
|
|
)
|
|
overlay(
|
|
np.array(images[index])[:, :, ::-1].copy(),
|
|
[
|
|
[
|
|
[x["bbox"][0], x["bbox"][1]],
|
|
[x["bbox"][2], x["bbox"][1]],
|
|
[x["bbox"][2], x["bbox"][3]],
|
|
[x["bbox"][0], x["bbox"][3]],
|
|
]
|
|
for x in right
|
|
],
|
|
args.output_dir / f"layout-{index}.png",
|
|
)
|
|
report: dict[str, Any] = {
|
|
"layout": {
|
|
"single_seconds": single_seconds,
|
|
"batch_seconds": batch_seconds,
|
|
"single_calls": single_calls,
|
|
"batch_calls": observer.calls,
|
|
"checks": checks,
|
|
}
|
|
}
|
|
write_json(args.output_dir / "layout-single.json", single)
|
|
write_json(args.output_dir / "layout-batch.json", batch)
|
|
write_json(args.output_dir / "report.json", report)
|
|
for kind in ("ocr_det", "seal"):
|
|
seal = kind == "seal"
|
|
detector = TextDetectorONNX(
|
|
str(repo / "OCR/paddleocr" / ("seal_PP-OCRv4_det_infer.onnx" if seal else "ch_PP-OCRv6_tiny_det_infer.onnx")),
|
|
**(
|
|
{
|
|
"limit_side_len": 736,
|
|
"limit_type": "min",
|
|
"thresh": 0.2,
|
|
"box_thresh": 0.6,
|
|
"unclip_ratio": 0.5,
|
|
"box_type": "poly",
|
|
}
|
|
if seal
|
|
else {}
|
|
),
|
|
)
|
|
originals = [np.array(page)[:, :, ::-1].copy() for page in ([pages[2]] if seal else pages[:2])]
|
|
samples = [value for page in originals for value in (page, page[:, ::-1].copy())]
|
|
samples.insert(1, np.empty((0, 0, 3), np.uint8))
|
|
detector(originals[0])
|
|
observer = RecordingSession(detector.session, kind, retain_outputs=True)
|
|
detector.session = observer
|
|
started = time.perf_counter()
|
|
single = detector.batch_predict(samples, max_batch_size=1)
|
|
single_seconds = time.perf_counter() - started
|
|
single_calls = list(observer.calls)
|
|
raw_single = [output[0][0] for output in observer.outputs]
|
|
observer.calls.clear()
|
|
observer.outputs.clear()
|
|
started = time.perf_counter()
|
|
batch = detector.batch_predict(samples, max_batch_size=16)
|
|
batch_seconds = time.perf_counter() - started
|
|
# 本样本按尺寸连续排列；按桶首见顺序处理尾批后，可直接对照概率图。
|
|
raw_batch = [image for output in observer.outputs for image in output[0]]
|
|
checks = []
|
|
for index, (left, right) in enumerate(zip(single, batch)):
|
|
a, b = left[0], right[0]
|
|
equal = (
|
|
a is None and b is None
|
|
if a is None or b is None
|
|
else (len(a) == len(b) and all(np.array_equal(x, y) for x, y in zip(a, b)))
|
|
)
|
|
checks.append(
|
|
{
|
|
"index": index,
|
|
"boxes_match": bool(equal),
|
|
"single_count": 0 if a is None else len(a),
|
|
"batch_count": 0 if b is None else len(b),
|
|
}
|
|
)
|
|
if b is not None:
|
|
overlay(samples[index], b, args.output_dir / f"{kind}-{index}.png")
|
|
probability = [
|
|
{"max_abs": float(np.abs(a - b).max()), "allclose": bool(np.allclose(a, b, rtol=1e-4, atol=1e-5))}
|
|
for a, b in zip(raw_single, raw_batch)
|
|
]
|
|
report[kind] = {
|
|
"single_seconds": single_seconds,
|
|
"batch_seconds": batch_seconds,
|
|
"single_calls": single_calls,
|
|
"batch_calls": observer.calls,
|
|
"checks": checks,
|
|
"probabilities": probability,
|
|
}
|
|
write_json(args.output_dir / f"{kind}-single.json", [x[0] for x in single])
|
|
write_json(args.output_dir / f"{kind}-batch.json", [x[0] for x in batch])
|
|
write_json(args.output_dir / "report.json", report)
|
|
assert all(x["boxes_match"] for x in checks) and all(x["allclose"] for x in probability)
|
|
assert all(x["exact_match"] for x in report["layout"]["checks"])
|
|
return report
|
|
|
|
|
|
def integration(args: argparse.Namespace) -> dict[str, Any]:
|
|
"""对真实 Parser 临时安装会话观察器，证明 PDF 调用链实际合批且不加载 Torch。"""
|
|
observers = []
|
|
|
|
def factory(model_path: str, device: str | None = None, intra_op_num_threads: int = 0) -> RecordingSession:
|
|
"""验证用会话工厂，生产会话配置原样透传。"""
|
|
observer = RecordingSession(ort_session(model_path, device, intra_op_num_threads), Path(model_path).name)
|
|
observers.append(observer)
|
|
return observer
|
|
|
|
with (
|
|
patch("mineru.model.layout.pp_doclayout_v2_onnx.ort_session", factory),
|
|
patch("mineru.model.ocr.pp_ocr_v6_onnx.ort_session", factory),
|
|
):
|
|
report = parse_documents(args)
|
|
report["sessions"] = [{"name": observer.name, "calls": observer.calls} for observer in observers]
|
|
layout_calls = [call for observer in observers if observer.name == "inference.onnx" for call in observer.calls]
|
|
det_calls = [call for observer in observers if "det_infer" in observer.name for call in observer.calls]
|
|
assert any(call["inputs"]["image"][0] > 1 for call in layout_calls)
|
|
assert any(next(iter(call["inputs"].values()))[0] > 1 for call in det_calls)
|
|
return report
|
|
|
|
|
|
def main() -> None:
|
|
"""显式选择组件对照或真实 PDF 集成验证，并保存配置与内存指标。"""
|
|
from mineru.config import config
|
|
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("kind", choices=["components", "parse"])
|
|
parser.add_argument("inputs", nargs="*", type=Path)
|
|
parser.add_argument("--models-dir", type=Path, default=Path("output/model-migration/models"))
|
|
parser.add_argument("--samples-dir", type=Path, default=Path("output/model-migration/samples"))
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--pages", default="")
|
|
parser.add_argument("--ocr-mode", default="auto", choices=["auto", "ocr", "txt"])
|
|
args = parser.parse_args()
|
|
args.small_backend, args.tier, args.rounds = "light", "basic", 1
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
config.model.base_dir, config.model.source, config.model.small_backend = str(args.models_dir.resolve()), "local", "onnx"
|
|
config.llm_aided.features.title_leveling = False
|
|
config.llm_aided.features.cross_page_table_cell_merge = False
|
|
report = components(args) if args.kind == "components" else integration(args)
|
|
report["peak_rss_bytes"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * (1 if sys.platform == "darwin" else 1024)
|
|
report["environment"] = {"python": sys.version.split()[0], "onnxruntime": version("onnxruntime")}
|
|
report["heavy_modules"] = [name for name in ("torch", "transformers") if name in sys.modules]
|
|
assert not report["heavy_modules"]
|
|
write_json(args.output_dir / "report.json", report)
|
|
print("Saved", args.output_dir, flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|