kuuhaku
/
PP-OCRv6_Demo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat:新增参数矩阵测试程序benchmark.py

This commit is contained in:
kuuhaku 2026-07-17 11:34:57 +08:00
parent 78bd690b7e
commit bf9f638695
5 changed files with 722 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -17,5 +17,7 @@ logs/*
!logs/.gitkeep
benchmarks/cpu/*.json
benchmarks/gpu/*.json
benchmarks/parameter-runs/
benchmarks/参数测试报告-*.md
!benchmarks/cpu/.gitkeep
!benchmarks/gpu/.gitkeep

40
README.md
View File

@ -187,10 +187,48 @@ outputs/
JSON 中保留坐标，调用方可以按具体业务继续处理。
## 参数速度测试
使用根目录 `benchmark.py` 对 `data/` 中的图片和 PDF 执行参数矩阵测试，并生成 Markdown 对比报告：
```bash
# 默认测试 tiny/small/medium × fast/standard，PDF 强制 OCR
python benchmark.py
# 显式指定模型、预处理配置、DPI 和 CPU 线程
python benchmark.py \
--models tiny small medium \
--profiles fast standard robust \
--dpis 120 144 200 \
--threads auto 4 8 \
--det-limit-side-lens 64 960 \
--det-thresholds 0.2 0.3 \
--det-box-thresholds 0.5 0.6 \
--det-unclip-ratios 1.5 2.0 \
--rec-score-thresholds 0.0 0.5 \
--rec-batch-sizes 1 6
# 指定报告路径
python benchmark.py --output benchmarks/参数测试报告.md
```
默认场景共 6 组：
- 三种模型：`tiny`、`small`、`medium`；
- 两种预处理配置：`fast`、`standard`；
- PDF 使用 `ocr` 模式，确保不同模型真正参与 PDF 测试；
- 每个场景在独立进程中执行，避免模型缓存到内存造成场景间干扰。
可测试的参数维度包括模型规格、预处理配置、PDF DPI、CPU 线程、检测边长及限制方式、检测阈值、文本框阈值、扩张比例、识别阈值和识别 Batch Size。所有列表参数会组成笛卡尔积。
报告统计墙钟耗时、模型初始化、图片推理、PDF OCR、纯 OCR 耗时、吞吐、识别行数和平均置信度。原始输出、日志及汇总 JSON 保存在 `benchmarks/parameter-runs/`。
> `robust` 会启用文档去畸变，测试显著更慢；多组 DPI 和线程参数会产生笛卡尔积，请控制测试规模。
## 测试
```bash
uv run --project cpu pytest -q
```
测试使用假模型覆盖结果适配、输出路由和 PDF 混合流程，不要求每次测试都下载 OCR 模型。
测试使用假模型覆盖结果适配、输出路由、参数测试程序和 PDF 混合流程，不要求每次测试都下载 OCR 模型。

584
benchmark.py Normal file
View File

@ -0,0 +1,584 @@
"""Run PP-OCRv6 parameter combinations against data/ and write a Markdown report."""
from __future__ import annotations
import argparse
import itertools
import json
import shutil
import statistics
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
ROOT = Path(__file__).resolve().parent
DEFAULT_DATA = ROOT / "data"
DEFAULT_BENCHMARK_ROOT = ROOT / "benchmarks" / "parameter-runs"
MODEL_SIZES = ("tiny", "small", "medium")
PROFILE_OPTIONS = {
"fast": {
"doc_orientation": False,
"doc_unwarping": False,
"textline_orientation": False,
"description": "关闭文档方向、去畸变和文本行方向，优先速度",
},
"standard": {
"doc_orientation": True,
"doc_unwarping": False,
"textline_orientation": True,
"description": "开启文档方向和文本行方向，关闭去畸变",
},
"robust": {
"doc_orientation": True,
"doc_unwarping": True,
"textline_orientation": True,
"description": "开启方向处理和文档去畸变，优先复杂输入适应性",
},
}
@dataclass(frozen=True)
class Scenario:
model_size: str
profile: str
dpi: int
threads: int | None
det_limit_side_len: int = 64
det_limit_type: str = "min"
det_thresh: float = 0.3
det_box_thresh: float = 0.6
det_unclip_ratio: float = 1.5
rec_score_thresh: float = 0.0
rec_batch_size: int = 6
@staticmethod
def _slug(value: float) -> str:
return f"{value:g}".replace("-", "m").replace(".", "p")
@property
def name(self) -> str:
thread_name = "auto" if self.threads is None else str(self.threads)
return (
f"{self.model_size}-{self.profile}-dpi{self.dpi}-th{thread_name}"
f"-s{self.det_limit_side_len}{self.det_limit_type}"
f"-d{self._slug(self.det_thresh)}-b{self._slug(self.det_box_thresh)}"
f"-u{self._slug(self.det_unclip_ratio)}-r{self._slug(self.rec_score_thresh)}"
f"-rb{self.rec_batch_size}"
)
@dataclass
class ScenarioResult:
scenario: Scenario
status: str
exit_code: int
wall_seconds: float
model_init_seconds: float | None
pure_ocr_seconds: float
image_inference_seconds: float
pdf_ocr_seconds: float
processed_files: int
image_files: int
pdf_files: int
ocr_pages: int
recognized_lines: int
mean_confidence: float | None
units_per_second: float | None
output_dir: str
log_file: str
command: list[str]
error: str | None = None
def parse_threads(values: list[str]) -> list[int | None]:
parsed: list[int | None] = []
for value in values:
normalized = value.strip().lower()
if normalized == "auto":
item = None
else:
item = int(normalized)
if item < 1:
raise ValueError("线程数必须大于等于 1")
if item not in parsed:
parsed.append(item)
return parsed
def build_scenarios(
model_sizes: list[str],
profiles: list[str],
dpis: list[int],
threads: list[int | None],
det_limit_side_lens: list[int] | None = None,
det_limit_types: list[str] | None = None,
det_thresholds: list[float] | None = None,
det_box_thresholds: list[float] | None = None,
det_unclip_ratios: list[float] | None = None,
rec_score_thresholds: list[float] | None = None,
rec_batch_sizes: list[int] | None = None,
) -> list[Scenario]:
dimensions = itertools.product(
model_sizes,
profiles,
dpis,
threads,
det_limit_side_lens or [64],
det_limit_types or ["min"],
det_thresholds or [0.3],
det_box_thresholds or [0.6],
det_unclip_ratios or [1.5],
rec_score_thresholds or [0.0],
rec_batch_sizes or [6],
)
return [Scenario(*values) for values in dimensions]
def _bool_flag(name: str, enabled: bool) -> str:
return f"--{name}" if enabled else f"--no-{name}"
def build_command(
scenario: Scenario,
*,
data_dir: Path,
output_dir: Path,
device: str,
device_id: int,
warmup: int,
rounds: int,
pdf_mode: str,
) -> list[str]:
profile = PROFILE_OPTIONS[scenario.profile]
command = [
sys.executable,
str(ROOT / "ocr.py"),
str(data_dir),
"--recursive",
"--device",
device,
"--device-id",
str(device_id),
"--model-size",
scenario.model_size,
"--pdf-mode",
pdf_mode,
"--dpi",
str(scenario.dpi),
"--warmup",
str(warmup),
"--rounds",
str(rounds),
"--text-det-limit-side-len",
str(scenario.det_limit_side_len),
"--text-det-limit-type",
scenario.det_limit_type,
"--text-det-thresh",
str(scenario.det_thresh),
"--text-det-box-thresh",
str(scenario.det_box_thresh),
"--text-det-unclip-ratio",
str(scenario.det_unclip_ratio),
"--text-rec-score-thresh",
str(scenario.rec_score_thresh),
"--text-recognition-batch-size",
str(scenario.rec_batch_size),
"--output",
str(output_dir),
"--overwrite",
"--no-result",
_bool_flag("doc-orientation-classify", profile["doc_orientation"]),
_bool_flag("doc-unwarping", profile["doc_unwarping"]),
_bool_flag("textline-orientation", profile["textline_orientation"]),
]
if scenario.threads is not None:
command.extend(["--threads", str(scenario.threads)])
return command
def _read_json(path: Path) -> dict[str, Any]:
return json.loads(path.read_text(encoding="utf-8"))
def _collect_confidences(payload: dict[str, Any]) -> list[float]:
scores: list[float] = []
for line in payload.get("lines", []):
value = line.get("score")
if isinstance(value, (int, float)):
scores.append(float(value))
return scores
def collect_scenario_metrics(
scenario: Scenario,
*,
output_dir: Path,
log_file: Path,
command: list[str],
exit_code: int,
wall_seconds: float,
error: str | None,
) -> ScenarioResult:
batch_manifests = sorted((output_dir / "batches").glob("*.json"))
batch = _read_json(batch_manifests[-1]) if batch_manifests else {}
image_benchmarks = list((output_dir / "images").rglob("benchmark.json"))
image_results = list((output_dir / "images").rglob("result.json"))
pdf_manifests = list((output_dir / "pdfs").rglob("manifest.json"))
pdf_page_results = list((output_dir / "pdfs").rglob("pages/page-*.json"))
image_inference = 0.0
model_init_values: list[float] = []
recognized_lines = 0
confidences: list[float] = []
for path in image_benchmarks:
payload = _read_json(path)
image_inference += sum(float(value) for value in payload.get("inference_seconds", {}).get("all", []))
value = payload.get("model_init_seconds")
if isinstance(value, (int, float)):
model_init_values.append(float(value))
for path in image_results:
payload = _read_json(path)
recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0))
confidences.extend(_collect_confidences(payload))
pdf_ocr_seconds = 0.0
ocr_pages = 0
for path in pdf_manifests:
payload = _read_json(path)
summary = payload.get("summary", {})
ocr_pages += int(summary.get("ocr_pages", 0))
metadata = payload.get("run_metadata", {})
value = metadata.get("model_init_seconds")
if isinstance(value, (int, float)):
model_init_values.append(float(value))
for page in payload.get("pages", {}).values():
pdf_ocr_seconds += float(page.get("ocr_seconds", 0.0) or 0.0)
if page.get("source_type") == "text":
recognized_lines += int(page.get("detected_lines", 0) or 0)
for path in pdf_page_results:
payload = _read_json(path)
if "lines" in payload:
recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0))
confidences.extend(_collect_confidences(payload))
pure_ocr_seconds = image_inference + pdf_ocr_seconds
image_files = len(image_results)
pdf_files = len(pdf_manifests)
units = image_files + ocr_pages
units_per_second = units / pure_ocr_seconds if pure_ocr_seconds > 0 else None
status = "completed" if exit_code == 0 else "failed"
if exit_code == 0 and not batch:
status = "incomplete"
error = error or "未找到批处理 manifest"
return ScenarioResult(
scenario=scenario,
status=status,
exit_code=exit_code,
wall_seconds=wall_seconds,
model_init_seconds=max(model_init_values) if model_init_values else None,
pure_ocr_seconds=pure_ocr_seconds,
image_inference_seconds=image_inference,
pdf_ocr_seconds=pdf_ocr_seconds,
processed_files=int(batch.get("completed_files", image_files + pdf_files)),
image_files=image_files,
pdf_files=pdf_files,
ocr_pages=ocr_pages,
recognized_lines=recognized_lines,
mean_confidence=statistics.fmean(confidences) if confidences else None,
units_per_second=units_per_second,
output_dir=str(output_dir),
log_file=str(log_file),
command=command,
error=error,
)
def _format_seconds(value: float | None) -> str:
return "-" if value is None else f"{value:.3f}"
def _format_float(value: float | None, digits: int = 4) -> str:
return "-" if value is None else f"{value:.{digits}f}"
def render_report(
results: list[ScenarioResult],
*,
data_dir: Path,
device: str,
pdf_mode: str,
warmup: int,
rounds: int,
started_at: datetime,
finished_at: datetime,
) -> str:
completed = [result for result in results if result.status == "completed"]
by_ocr = sorted(completed, key=lambda item: item.pure_ocr_seconds or float("inf"))
by_wall = sorted(completed, key=lambda item: item.wall_seconds)
lines = [
"# PP-OCRv6 参数测试报告",
"",
"## 测试信息",
"",
f"- 数据目录：`{data_dir}`",
f"- 运行设备：`{device}`",
f"- PDF 模式：`{pdf_mode}`",
f"- 图片预热轮数：`{warmup}`",
f"- 图片推理轮数：`{rounds}`",
f"- 场景数量：`{len(results)}`",
f"- 开始时间：`{started_at.astimezone().isoformat()}`",
f"- 结束时间：`{finished_at.astimezone().isoformat()}`",
f"- 总耗时：`{(finished_at - started_at).total_seconds():.3f}s`",
"",
"> 纯 OCR 耗时为图片推理耗时与 PDF OCR 页面耗时之和，不含模型初始化、文件扫描、PDF 文本层提取和结果导出。首次下载模型会显著增加墙钟耗时，应优先参考缓存模型后的结果。",
"",
"## 汇总对比",
"",
"| 排名 | 场景 | 状态 | 墙钟耗时(s) | 模型初始化(s) | 纯OCR耗时(s) | 图片推理(s) | PDF OCR(s) | OCR单位/秒 | 识别行数 | 平均置信度 |",
"|---:|---|---|---:|---:|---:|---:|---:|---:|---:|---:|",
]
rank_map = {result.scenario.name: index + 1 for index, result in enumerate(by_ocr)}
for result in results:
lines.append(
"| {rank} | `{name}` | {status} | {wall} | {init} | {pure} | {image} | {pdf} | {rate} | {count} | {confidence} |".format(
rank=rank_map.get(result.scenario.name, "-"),
name=result.scenario.name,
status=result.status,
wall=_format_seconds(result.wall_seconds),
init=_format_seconds(result.model_init_seconds),
pure=_format_seconds(result.pure_ocr_seconds),
image=_format_seconds(result.image_inference_seconds),
pdf=_format_seconds(result.pdf_ocr_seconds),
rate=_format_float(result.units_per_second, 3),
count=result.recognized_lines,
confidence=_format_float(result.mean_confidence),
)
)
lines.extend(["", "## 结论", ""])
if by_ocr:
fastest = by_ocr[0]
lines.append(
f"- 纯 OCR 耗时最短：`{fastest.scenario.name}`，耗时 `{fastest.pure_ocr_seconds:.3f}s`，吞吐 `{_format_float(fastest.units_per_second, 3)}` OCR 单位/秒。"
)
if by_wall:
fastest_wall = by_wall[0]
lines.append(
f"- 墙钟耗时最短：`{fastest_wall.scenario.name}`，耗时 `{fastest_wall.wall_seconds:.3f}s`。"
)
if len(by_ocr) > 1 and by_ocr[-1].pure_ocr_seconds > 0:
speedup = by_ocr[-1].pure_ocr_seconds / by_ocr[0].pure_ocr_seconds
lines.append(f"- 最快与最慢完成场景的纯 OCR 速度差约为 `{speedup:.2f}x`。")
lines.append("- 平均置信度仅作为结果稳定性的辅助观察值，不等价于有标注数据集上的识别准确率。")
lines.extend(["", "## 参数说明", ""])
for name, profile in PROFILE_OPTIONS.items():
lines.append(f"- `{name}`：{profile['description']}。")
lines.extend(
[
"",
"## 场景明细",
"",
]
)
for result in results:
scenario = result.scenario
lines.extend(
[
f"### {scenario.name}",
"",
f"- 状态：`{result.status}`；退出码：`{result.exit_code}`",
f"- 模型规格：`{scenario.model_size}`",
f"- 预处理配置：`{scenario.profile}`",
f"- PDF DPI：`{scenario.dpi}`",
f"- CPU 线程：`{'auto' if scenario.threads is None else scenario.threads}`",
f"- 检测边长：`{scenario.det_limit_side_len}`；限制方式：`{scenario.det_limit_type}`",
f"- 检测阈值：`{scenario.det_thresh}`；文本框阈值：`{scenario.det_box_thresh}`；扩张比例：`{scenario.det_unclip_ratio}`",
f"- 识别阈值：`{scenario.rec_score_thresh}`；识别 Batch Size：`{scenario.rec_batch_size}`",
f"- 完成文件：`{result.processed_files}`（图片 `{result.image_files}`，PDF `{result.pdf_files}`，OCR 页 `{result.ocr_pages}`）",
f"- 墙钟耗时：`{result.wall_seconds:.3f}s`",
f"- 纯 OCR 耗时：`{result.pure_ocr_seconds:.3f}s`",
f"- 输出目录：`{result.output_dir}`",
f"- 运行日志：`{result.log_file}`",
f"- 命令：`{' '.join(result.command)}`",
]
)
if result.error:
lines.append(f"- 错误：`{result.error}`")
lines.append("")
return "\n".join(lines).rstrip() + "\n"
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="对 data 目录执行 PP-OCRv6 参数矩阵测试并生成 Markdown 报告",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("--data", type=Path, default=DEFAULT_DATA, help="测试数据目录")
parser.add_argument("--device", choices=("cpu", "gpu"), default="cpu", help="测试设备")
parser.add_argument("--device-id", type=int, default=0, help="GPU 编号")
parser.add_argument("--models", nargs="+", choices=MODEL_SIZES, default=list(MODEL_SIZES), help="模型规格列表")
parser.add_argument("--profiles", nargs="+", choices=tuple(PROFILE_OPTIONS), default=["fast", "standard"], help="预处理配置列表")
parser.add_argument("--dpis", nargs="+", type=int, default=[144], help="PDF 渲染 DPI 列表")
parser.add_argument("--threads", nargs="+", default=["auto"], help="CPU 线程列表，可使用 auto")
parser.add_argument("--det-limit-side-lens", nargs="+", type=int, default=[64], help="文本检测边长列表")
parser.add_argument("--det-limit-types", nargs="+", choices=("min", "max"), default=["min"], help="文本检测边长限制方式列表")
parser.add_argument("--det-thresholds", nargs="+", type=float, default=[0.3], help="文本检测像素阈值列表")
parser.add_argument("--det-box-thresholds", nargs="+", type=float, default=[0.6], help="文本框阈值列表")
parser.add_argument("--det-unclip-ratios", nargs="+", type=float, default=[1.5], help="文本框扩张比例列表")
parser.add_argument("--rec-score-thresholds", nargs="+", type=float, default=[0.0], help="识别置信度阈值列表")
parser.add_argument("--rec-batch-sizes", nargs="+", type=int, default=[6], help="文本识别 Batch Size 列表")
parser.add_argument("--warmup", type=int, default=1, help="每个场景首张图片预热轮数")
parser.add_argument("--rounds", type=int, default=1, help="每张图片推理轮数")
parser.add_argument("--pdf-mode", choices=("ocr", "hybrid", "text"), default="ocr", help="PDF 测试模式；速度对比推荐 ocr")
parser.add_argument("--output", type=Path, default=None, help="Markdown 报告路径")
parser.add_argument("--work-dir", type=Path, default=None, help="各场景原始输出目录")
parser.add_argument("--overwrite", action="store_true", help="允许删除已存在的 work-dir")
parser.add_argument("--fail-fast", action="store_true", help="任一场景失败后停止")
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
data_dir = args.data.expanduser().resolve()
if not data_dir.is_dir():
print(f"测试数据目录不存在: {data_dir}", file=sys.stderr)
return 2
if args.warmup < 0 or args.rounds < 1:
print("--warmup 必须 >= 0，--rounds 必须 >= 1", file=sys.stderr)
return 2
if any(dpi < 72 or dpi > 600 for dpi in args.dpis):
print("--dpis 必须在 72 到 600 之间", file=sys.stderr)
return 2
if any(value < 1 for value in args.det_limit_side_lens):
print("--det-limit-side-lens 必须大于等于 1", file=sys.stderr)
return 2
if any(value < 1 for value in args.rec_batch_sizes):
print("--rec-batch-sizes 必须大于等于 1", file=sys.stderr)
return 2
try:
thread_values = parse_threads(args.threads)
except ValueError as exc:
print(f"无效线程参数: {exc}", file=sys.stderr)
return 2
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
work_dir = (args.work_dir or DEFAULT_BENCHMARK_ROOT / timestamp).expanduser().resolve()
report_path = (args.output or ROOT / "benchmarks" / f"参数测试报告-{timestamp}.md").expanduser().resolve()
if work_dir.exists():
if not args.overwrite:
print(f"测试输出目录已存在: {work_dir}；请使用 --overwrite", file=sys.stderr)
return 2
shutil.rmtree(work_dir)
work_dir.mkdir(parents=True, exist_ok=True)
report_path.parent.mkdir(parents=True, exist_ok=True)
scenarios = build_scenarios(
args.models,
args.profiles,
args.dpis,
thread_values,
args.det_limit_side_lens,
args.det_limit_types,
args.det_thresholds,
args.det_box_thresholds,
args.det_unclip_ratios,
args.rec_score_thresholds,
args.rec_batch_sizes,
)
print(f"计划执行 {len(scenarios)} 个参数场景，报告将写入: {report_path}")
started_at = datetime.now().astimezone()
results: list[ScenarioResult] = []
for index, scenario in enumerate(scenarios, 1):
scenario_dir = work_dir / scenario.name
output_dir = scenario_dir / "outputs"
log_file = scenario_dir / "run.log"
scenario_dir.mkdir(parents=True, exist_ok=True)
command = build_command(
scenario,
data_dir=data_dir,
output_dir=output_dir,
device=args.device,
device_id=args.device_id,
warmup=args.warmup,
rounds=args.rounds,
pdf_mode=args.pdf_mode,
)
print(f"[{index}/{len(scenarios)}] {scenario.name}")
wall_started = time.perf_counter()
error = None
try:
completed = subprocess.run(
command,
cwd=ROOT,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
check=False,
)
raw_output = completed.stdout or b""
log_file.write_bytes(raw_output)
exit_code = completed.returncode
if exit_code != 0:
text = raw_output.decode("utf-8", errors="replace")
error = next((line for line in reversed(text.splitlines()) if line.strip()), f"退出码 {exit_code}")
except Exception as exc:
exit_code = 1
error = f"{type(exc).__name__}: {exc}"
log_file.write_text(error + "\n", encoding="utf-8")
wall_seconds = time.perf_counter() - wall_started
result = collect_scenario_metrics(
scenario,
output_dir=output_dir,
log_file=log_file,
command=command,
exit_code=exit_code,
wall_seconds=wall_seconds,
error=error,
)
results.append(result)
print(
f" 状态={result.status} 墙钟={result.wall_seconds:.3f}s "
f"纯OCR={result.pure_ocr_seconds:.3f}s"
)
if result.status != "completed" and args.fail_fast:
break
finished_at = datetime.now().astimezone()
report = render_report(
results,
data_dir=data_dir,
device=args.device,
pdf_mode=args.pdf_mode,
warmup=args.warmup,
rounds=args.rounds,
started_at=started_at,
finished_at=finished_at,
)
report_path.write_text(report, encoding="utf-8")
summary_path = work_dir / "summary.json"
summary_path.write_text(
json.dumps(
[
{
**asdict(result),
"scenario": asdict(result.scenario),
}
for result in results
],
ensure_ascii=False,
indent=2,
),
encoding="utf-8",
)
print(f"测试完成，Markdown 报告: {report_path}")
print(f"原始汇总 JSON: {summary_path}")
return 0 if all(result.status == "completed" for result in results) else 3
if __name__ == "__main__":
raise SystemExit(main())

18
benchmarks/README.md
View File

@ -17,3 +17,21 @@ python ocr.py data/images/手写01.png --warmup 1 --rounds 3
```bash
python ocr.py data/images/手写01.png --benchmark-json benchmarks/cpu/手写01.json
```
## 参数矩阵测试
```bash
python benchmark.py
```
默认比较 `tiny/small/medium` 与 `fast/standard` 预处理配置，报告生成到：
```text
benchmarks/参数测试报告-<时间戳>.md
```
各场景原始结果和汇总 JSON 保存在：
```text
benchmarks/parameter-runs/<时间戳>/
```

79
tests/test_benchmark.py Normal file
View File

@ -0,0 +1,79 @@
from pathlib import Path
from benchmark import Scenario, build_command, build_scenarios, parse_threads, render_report, ScenarioResult
def test_build_default_matrix():
scenarios = build_scenarios(
["tiny", "small", "medium"],
["fast", "standard"],
[144],
[None],
)
assert len(scenarios) == 6
assert scenarios[0].name == "tiny-fast-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6"
assert scenarios[-1].name == "medium-standard-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6"
def test_parse_threads():
assert parse_threads(["auto", "4", "4"]) == [None, 4]
def test_command_contains_model_and_profile_flags(tmp_path):
scenario = Scenario("small", "fast", 120, 4)
command = build_command(
scenario,
data_dir=tmp_path / "data",
output_dir=tmp_path / "output",
device="cpu",
device_id=0,
warmup=1,
rounds=2,
pdf_mode="ocr",
)
assert command[1].endswith("ocr.py")
assert command[command.index("--model-size") + 1] == "small"
assert "--no-doc-orientation-classify" in command
assert "--no-textline-orientation" in command
assert command[command.index("--threads") + 1] == "4"
assert command[command.index("--text-det-thresh") + 1] == "0.3"
assert command[command.index("--text-recognition-batch-size") + 1] == "6"
def test_report_contains_comparison_table(tmp_path):
scenario = Scenario("tiny", "fast", 144, None)
result = ScenarioResult(
scenario=scenario,
status="completed",
exit_code=0,
wall_seconds=3.0,
model_init_seconds=1.0,
pure_ocr_seconds=2.0,
image_inference_seconds=1.0,
pdf_ocr_seconds=1.0,
processed_files=2,
image_files=1,
pdf_files=1,
ocr_pages=1,
recognized_lines=10,
mean_confidence=0.95,
units_per_second=1.0,
output_dir=str(tmp_path / "output"),
log_file=str(tmp_path / "run.log"),
command=["python", "ocr.py"],
)
from datetime import datetime
report = render_report(
[result],
data_dir=Path("data"),
device="cpu",
pdf_mode="ocr",
warmup=1,
rounds=1,
started_at=datetime.now().astimezone(),
finished_at=datetime.now().astimezone(),
)
assert "# PP-OCRv6 参数测试报告" in report
assert "tiny-fast-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6" in report
assert "纯 OCR 耗时最短" in report